From dae3e24175e6a0714686100ad1b745073d06106f Mon Sep 17 00:00:00 2001 From: Linyu <94553312+weijinglin@users.noreply.github.com> Date: Tue, 16 Sep 2025 14:54:03 +0800 Subject: [PATCH 01/71] Refactor: Refactor Scheduler to Support Dynamic Workflow Scheduling and Pipeline Pooling (#48) --- hugegraph-llm/pyproject.toml | 2 + .../src/hugegraph_llm/flows/__init__.py | 16 ++ .../hugegraph_llm/flows/build_vector_index.py | 55 ++++ .../src/hugegraph_llm/flows/common.py | 45 +++ .../src/hugegraph_llm/flows/graph_extract.py | 127 +++++++++ .../src/hugegraph_llm/flows/scheduler.py | 90 ++++++ .../models/embeddings/init_embedding.py | 36 ++- .../src/hugegraph_llm/models/llms/init_llm.py | 80 +++++- .../operators/common_op/check_schema.py | 258 ++++++++++++++++-- .../operators/document_op/chunk_split.py | 59 ++++ .../operators/hugegraph_op/schema_manager.py | 88 +++++- .../operators/index_op/build_vector_index.py | 65 ++++- .../operators/llm_op/info_extract.py | 220 +++++++++++++-- .../llm_op/property_graph_extract.py | 190 +++++++++++-- .../src/hugegraph_llm/operators/util.py | 27 ++ .../src/hugegraph_llm/state/__init__.py | 16 ++ .../src/hugegraph_llm/state/ai_state.py | 81 ++++++ .../hugegraph_llm/utils/graph_index_utils.py | 83 ++++-- .../hugegraph_llm/utils/vector_index_utils.py | 67 +++-- 19 files changed, 1472 insertions(+), 133 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/__init__.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/common.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/scheduler.py create mode 100644 hugegraph-llm/src/hugegraph_llm/operators/util.py create mode 100644 hugegraph-llm/src/hugegraph_llm/state/__init__.py create mode 100644 hugegraph-llm/src/hugegraph_llm/state/ai_state.py diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index 2ed438967..09c49ae26 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -58,6 +58,7 @@ dependencies = [ "apscheduler", "litellm", "hugegraph-python-client", + "pycgraph", ] [project.urls] homepage = "https://hugegraph.apache.org/" @@ -85,3 +86,4 @@ allow-direct-references = true [tool.uv.sources] hugegraph-python-client = { workspace = true } +pycgraph = { git = "https://github.com/ChunelFeng/CGraph.git", subdirectory = "python", rev = "main", marker = "sys_platform == 'linux'" } diff --git a/hugegraph-llm/src/hugegraph_llm/flows/__init__.py b/hugegraph-llm/src/hugegraph_llm/flows/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/__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/flows/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py new file mode 100644 index 000000000..f1ee8c1c4 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py @@ -0,0 +1,55 @@ +# 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.common import BaseFlow +from hugegraph_llm.state.ai_state import WkFlowInput + +import json +from PyCGraph import GPipeline + +from hugegraph_llm.operators.document_op.chunk_split import ChunkSplitNode +from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndexNode +from hugegraph_llm.state.ai_state import WkFlowState + + +class BuildVectorIndexFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, texts): + prepared_input.texts = texts + prepared_input.language = "zh" + prepared_input.split_type = "paragraph" + return + + def build_flow(self, texts): + pipeline = GPipeline() + # prepare for workflow input + prepared_input = WkFlowInput() + self.prepare(prepared_input, texts) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + chunk_split_node = ChunkSplitNode() + build_vector_node = BuildVectorIndexNode() + pipeline.registerGElement(chunk_split_node, set(), "chunk_split") + pipeline.registerGElement(build_vector_node, {chunk_split_node}, "build_vector") + + return pipeline + + def post_deal(self, pipeline=None): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + return json.dumps(res, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/common.py b/hugegraph-llm/src/hugegraph_llm/flows/common.py new file mode 100644 index 000000000..4c552626a --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/common.py @@ -0,0 +1,45 @@ +# 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 abc import ABC, abstractmethod + +from hugegraph_llm.state.ai_state import WkFlowInput + + +class BaseFlow(ABC): + """ + Base class for flows, defines three interface methods: prepare, build_flow, and post_deal. + """ + + @abstractmethod + def prepare(self, prepared_input: WkFlowInput, *args, **kwargs): + """ + Pre-processing interface. + """ + pass + + @abstractmethod + def build_flow(self, *args, **kwargs): + """ + Interface for building the flow. + """ + pass + + @abstractmethod + def post_deal(self, *args, **kwargs): + """ + Post-processing interface. + """ + pass diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py new file mode 100644 index 000000000..f1a6c5f6f --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py @@ -0,0 +1,127 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from PyCGraph import GPipeline +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.operators.common_op.check_schema import CheckSchemaNode +from hugegraph_llm.operators.document_op.chunk_split import ChunkSplitNode +from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManagerNode +from hugegraph_llm.operators.llm_op.info_extract import InfoExtractNode +from hugegraph_llm.operators.llm_op.property_graph_extract import ( + PropertyGraphExtractNode, +) +from hugegraph_llm.utils.log import log + + +class GraphExtractFlow(BaseFlow): + def __init__(self): + pass + + def _import_schema( + self, + from_hugegraph=None, + from_extraction=None, + from_user_defined=None, + ): + if from_hugegraph: + return SchemaManagerNode() + elif from_user_defined: + return CheckSchemaNode() + elif from_extraction: + raise NotImplementedError("Not implemented yet") + else: + raise ValueError("No input data / invalid schema type") + + def prepare( + self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type + ): + # prepare input data + prepared_input.texts = texts + prepared_input.language = "zh" + prepared_input.split_type = "document" + prepared_input.example_prompt = example_prompt + prepared_input.schema = schema + schema = schema.strip() + if schema.startswith("{"): + try: + schema = json.loads(schema) + prepared_input.schema = schema + except json.JSONDecodeError as exc: + log.error("Invalid JSON format in schema. Please check it again.") + raise ValueError("Invalid JSON format in schema.") from exc + else: + log.info("Get schema '%s' from graphdb.", schema) + prepared_input.graph_name = schema + return + + def build_flow(self, schema, texts, example_prompt, extract_type): + pipeline = GPipeline() + prepared_input = WkFlowInput() + # prepare input data + self.prepare(prepared_input, schema, texts, example_prompt, extract_type) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + schema = schema.strip() + schema_node = None + if schema.startswith("{"): + try: + schema = json.loads(schema) + schema_node = self._import_schema(from_user_defined=schema) + except json.JSONDecodeError as exc: + log.error("Invalid JSON format in schema. Please check it again.") + raise ValueError("Invalid JSON format in schema.") from exc + else: + log.info("Get schema '%s' from graphdb.", schema) + schema_node = self._import_schema(from_hugegraph=schema) + + chunk_split_node = ChunkSplitNode() + graph_extract_node = None + if extract_type == "triples": + graph_extract_node = InfoExtractNode() + elif extract_type == "property_graph": + graph_extract_node = PropertyGraphExtractNode() + else: + raise ValueError(f"Unsupported extract_type: {extract_type}") + pipeline.registerGElement(schema_node, set(), "schema_node") + pipeline.registerGElement(chunk_split_node, set(), "chunk_split") + pipeline.registerGElement( + graph_extract_node, {schema_node, chunk_split_node}, "graph_extract" + ) + + return pipeline + + def post_deal(self, pipeline=None): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + vertices = res.get("vertices", []) + edges = res.get("edges", []) + if not vertices and not edges: + log.info("Please check the schema.(The schema may not match the Doc)") + return json.dumps( + { + "vertices": vertices, + "edges": edges, + "warning": "The schema may not match the Doc", + }, + ensure_ascii=False, + indent=2, + ) + return json.dumps( + {"vertices": vertices, "edges": edges}, + ensure_ascii=False, + indent=2, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py new file mode 100644 index 000000000..b096310db --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -0,0 +1,90 @@ +# 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 threading +from typing import Dict, Any +from PyCGraph import GPipelineManager +from hugegraph_llm.flows.build_vector_index import BuildVectorIndexFlow +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.flows.graph_extract import GraphExtractFlow +from hugegraph_llm.utils.log import log + + +class Scheduler: + pipeline_pool: Dict[str, Any] = None + max_pipeline: int + + def __init__(self, max_pipeline: int = 10): + self.pipeline_pool = {} + # pipeline_pool act as a manager of GPipelineManager which used for pipeline management + self.pipeline_pool["build_vector_index"] = { + "manager": GPipelineManager(), + "flow": BuildVectorIndexFlow(), + } + self.pipeline_pool["graph_extract"] = { + "manager": GPipelineManager(), + "flow": GraphExtractFlow(), + } + self.max_pipeline = max_pipeline + + # TODO: Implement Agentic Workflow + def agentic_flow(self): + pass + + def schedule_flow(self, flow: str, *args, **kwargs): + if flow not in self.pipeline_pool: + raise ValueError(f"Unsupported workflow {flow}") + manager = self.pipeline_pool[flow]["manager"] + flow: BaseFlow = self.pipeline_pool[flow]["flow"] + pipeline = manager.fetch() + if pipeline is None: + # call coresponding flow_func to create new workflow + pipeline = flow.build_flow(*args, **kwargs) + status = pipeline.init() + if status.isErr(): + error_msg = f"Error in flow init: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + status = pipeline.run() + if status.isErr(): + error_msg = f"Error in flow execution: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + res = flow.post_deal(pipeline) + manager.add(pipeline) + return res + else: + # fetch pipeline & prepare input for flow + prepared_input = pipeline.getGParamWithNoEmpty("wkflow_input") + flow.prepare(prepared_input, *args, **kwargs) + status = pipeline.run() + if status.isErr(): + raise RuntimeError(f"Error in flow execution {status.getInfo()}") + res = flow.post_deal(pipeline) + manager.release(pipeline) + return res + + +class SchedulerSingleton: + _instance = None + _instance_lock = threading.Lock() + + @classmethod + def get_instance(cls): + if cls._instance is None: + with cls._instance_lock: + if cls._instance is None: + cls._instance = Scheduler() + return cls._instance diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index 48e4968c4..3ad50b3ec 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -17,10 +17,40 @@ from hugegraph_llm.config import llm_settings +from hugegraph_llm.config import LLMConfig from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding from hugegraph_llm.models.embeddings.openai import OpenAIEmbedding +model_map = { + "openai": llm_settings.openai_embedding_model, + "ollama/local": llm_settings.ollama_embedding_model, + "litellm": llm_settings.litellm_embedding_model, +} + + +def get_embedding(llm_settings: LLMConfig): + if llm_settings.embedding_type == "openai": + return OpenAIEmbedding( + model_name=llm_settings.openai_embedding_model, + api_key=llm_settings.openai_embedding_api_key, + api_base=llm_settings.openai_embedding_api_base, + ) + if llm_settings.embedding_type == "ollama/local": + return OllamaEmbedding( + model_name=llm_settings.ollama_embedding_model, + host=llm_settings.ollama_embedding_host, + port=llm_settings.ollama_embedding_port, + ) + if llm_settings.embedding_type == "litellm": + return LiteLLMEmbedding( + model_name=llm_settings.litellm_embedding_model, + api_key=llm_settings.litellm_embedding_api_key, + api_base=llm_settings.litellm_embedding_api_base, + ) + + raise Exception("embedding type is not supported !") + class Embeddings: def __init__(self): @@ -31,19 +61,19 @@ def get_embedding(self): return OpenAIEmbedding( model_name=llm_settings.openai_embedding_model, api_key=llm_settings.openai_embedding_api_key, - api_base=llm_settings.openai_embedding_api_base + api_base=llm_settings.openai_embedding_api_base, ) if self.embedding_type == "ollama/local": return OllamaEmbedding( model_name=llm_settings.ollama_embedding_model, host=llm_settings.ollama_embedding_host, - port=llm_settings.ollama_embedding_port + port=llm_settings.ollama_embedding_port, ) if self.embedding_type == "litellm": return LiteLLMEmbedding( model_name=llm_settings.litellm_embedding_model, api_key=llm_settings.litellm_embedding_api_key, - api_base=llm_settings.litellm_embedding_api_base + api_base=llm_settings.litellm_embedding_api_base, ) raise Exception("embedding type is not supported !") 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 e70b0d9d7..7e1eaab68 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py @@ -15,13 +15,85 @@ # specific language governing permissions and limitations # under the License. - +from hugegraph_llm.config import LLMConfig from hugegraph_llm.models.llms.ollama import OllamaClient from hugegraph_llm.models.llms.openai import OpenAIClient from hugegraph_llm.models.llms.litellm import LiteLLMClient from hugegraph_llm.config import llm_settings +def get_chat_llm(llm_settings: LLMConfig): + if llm_settings.chat_llm_type == "openai": + return OpenAIClient( + api_key=llm_settings.openai_chat_api_key, + api_base=llm_settings.openai_chat_api_base, + model_name=llm_settings.openai_chat_language_model, + max_tokens=llm_settings.openai_chat_tokens, + ) + if llm_settings.chat_llm_type == "ollama/local": + return OllamaClient( + model=llm_settings.ollama_chat_language_model, + host=llm_settings.ollama_chat_host, + port=llm_settings.ollama_chat_port, + ) + if llm_settings.chat_llm_type == "litellm": + return LiteLLMClient( + api_key=llm_settings.litellm_chat_api_key, + api_base=llm_settings.litellm_chat_api_base, + model_name=llm_settings.litellm_chat_language_model, + max_tokens=llm_settings.litellm_chat_tokens, + ) + raise Exception("chat llm type is not supported !") + + +def get_extract_llm(llm_settings: LLMConfig): + if llm_settings.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, + ) + if llm_settings.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, + ) + if llm_settings.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, + ) + raise Exception("extract llm type is not supported !") + + +def get_text2gql_llm(llm_settings: LLMConfig): + if llm_settings.text2gql_llm_type == "openai": + return OpenAIClient( + api_key=llm_settings.openai_text2gql_api_key, + api_base=llm_settings.openai_text2gql_api_base, + model_name=llm_settings.openai_text2gql_language_model, + max_tokens=llm_settings.openai_text2gql_tokens, + ) + if llm_settings.text2gql_llm_type == "ollama/local": + return OllamaClient( + model=llm_settings.ollama_text2gql_language_model, + host=llm_settings.ollama_text2gql_host, + port=llm_settings.ollama_text2gql_port, + ) + if llm_settings.text2gql_llm_type == "litellm": + return LiteLLMClient( + api_key=llm_settings.litellm_text2gql_api_key, + api_base=llm_settings.litellm_text2gql_api_base, + model_name=llm_settings.litellm_text2gql_language_model, + max_tokens=llm_settings.litellm_text2gql_tokens, + ) + raise Exception("text2gql llm type is not supported !") + + class LLMs: def __init__(self): self.chat_llm_type = llm_settings.chat_llm_type @@ -101,4 +173,8 @@ def get_text2gql_llm(self): if __name__ == "__main__": client = LLMs().get_chat_llm() print(client.generate(prompt="What is the capital of China?")) - print(client.generate(messages=[{"role": "user", "content": "What is the capital of China?"}])) + print( + client.generate( + messages=[{"role": "user", "content": "What is the capital of China?"}] + ) + ) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index 3220d9f3d..7a533517a 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -20,8 +20,12 @@ from hugegraph_llm.enums.property_cardinality import PropertyCardinality from hugegraph_llm.enums.property_data_type import PropertyDataType +from hugegraph_llm.operators.util import init_context from hugegraph_llm.utils.log import log +from PyCGraph import GNode, CStatus +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + def log_and_raise(message: str) -> None: log.warning(message) @@ -59,64 +63,270 @@ def _validate_schema(self, schema: Dict[str, Any]) -> None: check_type(schema, dict, "Input data is not a dictionary.") if "vertexlabels" not in schema or "edgelabels" not in schema: log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") - check_type(schema["vertexlabels"], list, "'vertexlabels' in input data is not a list.") - check_type(schema["edgelabels"], list, "'edgelabels' in input data is not a list.") + check_type( + schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." + ) + check_type( + schema["edgelabels"], list, "'edgelabels' in input data is not a list." + ) + + def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): + property_labels = schema.get("propertykeys", []) + check_type( + property_labels, + list, + "'propertykeys' in input data is not of correct type.", + ) + property_label_set = {label["name"] for label in property_labels} + return property_labels, property_label_set + + def _process_vertex_labels( + self, schema: Dict[str, Any], property_labels: list, property_label_set: set + ) -> None: + for vertex_label in schema["vertexlabels"]: + self._validate_vertex_label(vertex_label) + properties = vertex_label["properties"] + primary_keys = self._process_keys( + vertex_label, "primary_keys", properties[:1] + ) + if len(primary_keys) == 0: + log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") + vertex_label["primary_keys"] = primary_keys + nullable_keys = self._process_keys( + vertex_label, "nullable_keys", properties[1:] + ) + vertex_label["nullable_keys"] = nullable_keys + self._add_missing_properties( + properties, property_labels, property_label_set + ) + + def _process_edge_labels( + self, schema: Dict[str, Any], property_labels: list, property_label_set: set + ) -> None: + for edge_label in schema["edgelabels"]: + self._validate_edge_label(edge_label) + properties = edge_label.get("properties", []) + self._add_missing_properties( + properties, property_labels, property_label_set + ) + + def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: + check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") + if "name" not in vertex_label: + log_and_raise("VertexLabel in input data does not contain 'name'.") + check_type( + vertex_label["name"], str, "'name' in vertex_label is not of correct type." + ) + if "properties" not in vertex_label: + log_and_raise("VertexLabel in input data does not contain 'properties'.") + check_type( + vertex_label["properties"], + list, + "'properties' in vertex_label is not of correct type.", + ) + if len(vertex_label["properties"]) == 0: + log_and_raise("'properties' in vertex_label is empty.") + + def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: + check_type(edge_label, dict, "EdgeLabel in input data is not a dictionary.") + if ( + "name" not in edge_label + or "source_label" not in edge_label + or "target_label" not in edge_label + ): + log_and_raise( + "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." + ) + check_type( + edge_label["name"], str, "'name' in edge_label is not of correct type." + ) + check_type( + edge_label["source_label"], + str, + "'source_label' in edge_label is not of correct type.", + ) + check_type( + edge_label["target_label"], + str, + "'target_label' in edge_label is not of correct type.", + ) + + def _process_keys( + self, label: Dict[str, Any], key_type: str, default_keys: list + ) -> list: + keys = label.get(key_type, default_keys) + check_type( + keys, list, f"'{key_type}' in {label['name']} is not of correct type." + ) + new_keys = [key for key in keys if key in label["properties"]] + return new_keys + + def _add_missing_properties( + self, properties: list, property_labels: list, property_label_set: set + ) -> None: + for prop in properties: + if prop not in property_label_set: + property_labels.append( + { + "name": prop, + "data_type": PropertyDataType.DEFAULT.value, + "cardinality": PropertyCardinality.DEFAULT.value, + } + ) + property_label_set.add(prop) + + +class CheckSchemaNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + return init_context(self) + + def node_init(self): + if self.wk_input.schema is None: + return CStatus(-1, "Error occurs when prepare for workflow input") + self.data = self.wk_input.schema + return CStatus() + + def run(self) -> CStatus: + # init workflow input + sts = self.node_init() + if sts.isErr(): + return sts + # 1. Validate the schema structure + self.context.lock() + schema = self.data or self.context.schema + self._validate_schema(schema) + # 2. Process property labels and also create a set for it + property_labels, property_label_set = self._process_property_labels(schema) + # 3. Process properties in given vertex/edge labels + self._process_vertex_labels(schema, property_labels, property_label_set) + self._process_edge_labels(schema, property_labels, property_label_set) + # 4. Update schema with processed pks + schema["propertykeys"] = property_labels + self.context.schema = schema + self.context.unlock() + return CStatus() + + def _validate_schema(self, schema: Dict[str, Any]) -> None: + check_type(schema, dict, "Input data is not a dictionary.") + if "vertexlabels" not in schema or "edgelabels" not in schema: + log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") + check_type( + schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." + ) + check_type( + schema["edgelabels"], list, "'edgelabels' in input data is not a list." + ) def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): property_labels = schema.get("propertykeys", []) - check_type(property_labels, list, "'propertykeys' in input data is not of correct type.") + check_type( + property_labels, + list, + "'propertykeys' in input data is not of correct type.", + ) property_label_set = {label["name"] for label in property_labels} return property_labels, property_label_set - def _process_vertex_labels(self, schema: Dict[str, Any], property_labels: list, property_label_set: set) -> None: + def _process_vertex_labels( + self, schema: Dict[str, Any], property_labels: list, property_label_set: set + ) -> None: for vertex_label in schema["vertexlabels"]: self._validate_vertex_label(vertex_label) properties = vertex_label["properties"] - primary_keys = self._process_keys(vertex_label, "primary_keys", properties[:1]) + primary_keys = self._process_keys( + vertex_label, "primary_keys", properties[:1] + ) if len(primary_keys) == 0: log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") vertex_label["primary_keys"] = primary_keys - nullable_keys = self._process_keys(vertex_label, "nullable_keys", properties[1:]) + nullable_keys = self._process_keys( + vertex_label, "nullable_keys", properties[1:] + ) vertex_label["nullable_keys"] = nullable_keys - self._add_missing_properties(properties, property_labels, property_label_set) + self._add_missing_properties( + properties, property_labels, property_label_set + ) - def _process_edge_labels(self, schema: Dict[str, Any], property_labels: list, property_label_set: set) -> None: + def _process_edge_labels( + self, schema: Dict[str, Any], property_labels: list, property_label_set: set + ) -> None: for edge_label in schema["edgelabels"]: self._validate_edge_label(edge_label) properties = edge_label.get("properties", []) - self._add_missing_properties(properties, property_labels, property_label_set) + self._add_missing_properties( + properties, property_labels, property_label_set + ) def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") if "name" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'name'.") - check_type(vertex_label["name"], str, "'name' in vertex_label is not of correct type.") + check_type( + vertex_label["name"], str, "'name' in vertex_label is not of correct type." + ) if "properties" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'properties'.") - check_type(vertex_label["properties"], list, "'properties' in vertex_label is not of correct type.") + check_type( + vertex_label["properties"], + list, + "'properties' in vertex_label is not of correct type.", + ) if len(vertex_label["properties"]) == 0: log_and_raise("'properties' in vertex_label is empty.") def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: check_type(edge_label, dict, "EdgeLabel in input data is not a dictionary.") - if "name" not in edge_label or "source_label" not in edge_label or "target_label" not in edge_label: - log_and_raise("EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'.") - check_type(edge_label["name"], str, "'name' in edge_label is not of correct type.") - check_type(edge_label["source_label"], str, "'source_label' in edge_label is not of correct type.") - check_type(edge_label["target_label"], str, "'target_label' in edge_label is not of correct type.") + if ( + "name" not in edge_label + or "source_label" not in edge_label + or "target_label" not in edge_label + ): + log_and_raise( + "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." + ) + check_type( + edge_label["name"], str, "'name' in edge_label is not of correct type." + ) + check_type( + edge_label["source_label"], + str, + "'source_label' in edge_label is not of correct type.", + ) + check_type( + edge_label["target_label"], + str, + "'target_label' in edge_label is not of correct type.", + ) - def _process_keys(self, label: Dict[str, Any], key_type: str, default_keys: list) -> list: + def _process_keys( + self, label: Dict[str, Any], key_type: str, default_keys: list + ) -> list: keys = label.get(key_type, default_keys) - check_type(keys, list, f"'{key_type}' in {label['name']} is not of correct type.") + check_type( + keys, list, f"'{key_type}' in {label['name']} is not of correct type." + ) new_keys = [key for key in keys if key in label["properties"]] return new_keys - def _add_missing_properties(self, properties: list, property_labels: list, property_label_set: set) -> None: + def _add_missing_properties( + self, properties: list, property_labels: list, property_label_set: set + ) -> None: for prop in properties: if prop not in property_label_set: - property_labels.append({ - "name": prop, - "data_type": PropertyDataType.DEFAULT.value, - "cardinality": PropertyCardinality.DEFAULT.value, - }) + property_labels.append( + { + "name": prop, + "data_type": PropertyDataType.DEFAULT.value, + "cardinality": PropertyCardinality.DEFAULT.value, + } + ) property_label_set.add(prop) + + def get_result(self): + self.context.lock() + res = self.context.to_json() + self.context.unlock() + return res 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 8c2dd80f5..d779a40ab 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 @@ -19,6 +19,8 @@ from typing import Literal, Dict, Any, Optional, Union, List from langchain_text_splitters import RecursiveCharacterTextSplitter +from hugegraph_llm.operators.util import init_context +from PyCGraph import GNode, CStatus # Constants LANGUAGE_ZH = "zh" @@ -27,6 +29,63 @@ SPLIT_TYPE_PARAGRAPH = "paragraph" SPLIT_TYPE_SENTENCE = "sentence" + +class ChunkSplitNode(GNode): + def init(self): + return init_context(self) + + def node_init(self): + if ( + self.wk_input.texts is None + or self.wk_input.language is None + or self.wk_input.split_type is None + ): + return CStatus(-1, "Error occurs when prepare for workflow input") + texts = self.wk_input.texts + language = self.wk_input.language + split_type = self.wk_input.split_type + if isinstance(texts, str): + texts = [texts] + self.texts = texts + self.separators = self._get_separators(language) + self.text_splitter = self._get_text_splitter(split_type) + return CStatus() + + def _get_separators(self, language: str) -> List[str]: + if language == LANGUAGE_ZH: + return ["\n\n", "\n", "。", "，", ""] + if language == LANGUAGE_EN: + return ["\n\n", "\n", ".", ",", " ", ""] + raise ValueError("language must be zh or en") + + 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 + if split_type == SPLIT_TYPE_SENTENCE: + return RecursiveCharacterTextSplitter( + chunk_size=50, chunk_overlap=0, separators=self.separators + ).split_text + raise ValueError("Type must be document, paragraph or sentence") + + def run(self): + sts = self.node_init() + if sts.isErr(): + return sts + all_chunks = [] + for text in self.texts: + chunks = self.text_splitter(text) + all_chunks.extend(chunks) + + self.context.lock() + self.context.chunks = all_chunks + self.context.unlock() + return CStatus() + + class ChunkSplit: def __init__( self, 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 2f50bb818..670c18b4a 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 @@ -17,8 +17,12 @@ from typing import Dict, Any, Optional from hugegraph_llm.config import huge_settings +from hugegraph_llm.operators.util import init_context +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState from pyhugegraph.client import PyHugeClient +from PyCGraph import GNode, CStatus + class SchemaManager: def __init__(self, graph_name: str): @@ -39,15 +43,22 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} + new_vertex = { + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex + } mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) if "edgelabels" in schema: mini_schema["edgelabels"] = [] for edge in schema["edgelabels"]: - new_edge = {key: edge[key] for key in - ["name", "source_label", "target_label", "properties"] if key in edge} + new_edge = { + key: edge[key] + for key in ["name", "source_label", "target_label", "properties"] + if key in edge + } mini_schema["edgelabels"].append(new_edge) return mini_schema @@ -63,3 +74,74 @@ def run(self, context: Optional[Dict[str, Any]]) -> Dict[str, Any]: # TODO: enhance the logic here context["simple_schema"] = self.simple_schema(schema) return context + + +class SchemaManagerNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + return init_context(self) + + def node_init(self): + if self.wk_input.graph_name is None: + return CStatus(-1, "Error occurs when prepare for workflow input") + graph_name = self.wk_input.graph_name + self.graph_name = graph_name + self.client = PyHugeClient( + url=huge_settings.graph_url, + graph=self.graph_name, + user=huge_settings.graph_user, + pwd=huge_settings.graph_pwd, + graphspace=huge_settings.graph_space, + ) + self.schema = self.client.schema() + return CStatus() + + def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: + mini_schema = {} + + # Add necessary vertexlabels items (3) + if "vertexlabels" in schema: + mini_schema["vertexlabels"] = [] + for vertex in schema["vertexlabels"]: + new_vertex = { + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex + } + mini_schema["vertexlabels"].append(new_vertex) + + # Add necessary edgelabels items (4) + if "edgelabels" in schema: + mini_schema["edgelabels"] = [] + for edge in schema["edgelabels"]: + new_edge = { + key: edge[key] + for key in ["name", "source_label", "target_label", "properties"] + if key in edge + } + mini_schema["edgelabels"].append(new_edge) + + return mini_schema + + def run(self) -> CStatus: + sts = self.node_init() + if sts.isErr(): + return sts + schema = self.schema.getSchema() + if not schema["vertexlabels"] and not schema["edgelabels"]: + raise Exception(f"Can not get {self.graph_name}'s schema from HugeGraph!") + + self.context.lock() + self.context.schema = schema + # TODO: enhance the logic here + self.context.simple_schema = self.simple_schema(schema) + self.context.unlock() + return CStatus() + + def get_result(self): + self.context.lock() + res = self.context.to_json() + self.context.unlock() + return res diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py index ffb35564b..ee89d330f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py @@ -23,20 +23,75 @@ from hugegraph_llm.config import huge_settings, resource_path, llm_settings from hugegraph_llm.indices.vector_index import VectorIndex from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel, get_filename_prefix, get_index_folder_name +from hugegraph_llm.utils.embedding_utils import ( + get_embeddings_parallel, + get_filename_prefix, + get_index_folder_name, +) from hugegraph_llm.utils.log import log +from hugegraph_llm.operators.util import init_context +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from PyCGraph import GNode, CStatus + + +class BuildVectorIndexNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + return init_context(self) + + def node_init(self): + self.embedding = get_embedding(llm_settings) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(self.embedding, "model_name", None) + ) + self.vector_index = VectorIndex.from_index_file( + self.index_dir, self.filename_prefix + ) + return CStatus() + + def run(self): + # init workflow input + sts = self.node_init() + if sts.isErr(): + return sts + self.context.lock() + try: + if self.context.chunks is None: + raise ValueError("chunks not found in context.") + chunks = self.context.chunks + finally: + self.context.unlock() + chunks_embedding = [] + log.debug("Building vector index for %s chunks...", len(chunks)) + # TODO: use async_get_texts_embedding instead of single sync method + chunks_embedding = asyncio.run(get_embeddings_parallel(self.embedding, chunks)) + if len(chunks_embedding) > 0: + self.vector_index.add(chunks_embedding, chunks) + self.vector_index.to_index_file(self.index_dir, self.filename_prefix) + return CStatus() + class BuildVectorIndex: def __init__(self, embedding: BaseEmbedding): self.embedding = embedding - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, - getattr(self.embedding, "model_name", None) + llm_settings.embedding_type, getattr(self.embedding, "model_name", None) + ) + self.vector_index = VectorIndex.from_index_file( + self.index_dir, self.filename_prefix ) - self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: if "chunks" not in context: 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 42bb6b108..15a8fdda7 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 @@ -18,20 +18,26 @@ import re from typing import List, Any, Dict, Optional +from hugegraph_llm.config import llm_settings from hugegraph_llm.document.chunk_split import ChunkSplitter from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.utils.log import log +from hugegraph_llm.operators.util import init_context +from hugegraph_llm.models.llms.init_llm import get_chat_llm +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from PyCGraph import GNode, CStatus + SCHEMA_EXAMPLE_PROMPT = """## Main Task Extract Triples from the given text and graph schema ## Basic Rules 1. The output format must be: (X,Y,Z) - LABEL -In this format, Y must be a value from "properties" or "edge_label", +In this format, Y must be a value from "properties" or "edge_label", and LABEL must be X's vertex_label or Y's edge_label. 2. Don't extract attribute/property fields that do not exist in the given schema 3. Ensure the extract property is in the same type as the schema (like 'age' should be a number) -4. Translate the given schema filed into Chinese if the given text is Chinese but the schema is in English (Optional) +4. Translate the given schema filed into Chinese if the given text is Chinese but the schema is in English (Optional) ## Example (Note: Update the example to correspond to the given text and schema) ### Input example: @@ -75,8 +81,10 @@ def generate_extract_triple_prompt(text, schema=None) -> str: if schema: return schema_real_prompt - log.warning("Recommend to provide a graph schema to improve the extraction accuracy. " - "Now using the default schema.") + log.warning( + "Recommend to provide a graph schema to improve the extraction accuracy. " + "Now using the default schema." + ) return text_based_prompt @@ -105,11 +113,17 @@ def extract_triples_by_regex_with_schema(schema, text, graph): # TODO: use a more efficient way to compare the extract & input property p_lower = p.lower() for vertex in schema["vertices"]: - if vertex["vertex_label"] == label and any(pp.lower() == p_lower - for pp in vertex["properties"]): + if vertex["vertex_label"] == label and any( + pp.lower() == p_lower for pp in vertex["properties"] + ): id = f"{label}-{s}" if id not in vertices_dict: - vertices_dict[id] = {"id": id, "name": s, "label": label, "properties": {p: o}} + vertices_dict[id] = { + "id": id, + "name": s, + "label": label, + "properties": {p: o}, + } else: vertices_dict[id]["properties"].update({p: o}) break @@ -118,25 +132,35 @@ def extract_triples_by_regex_with_schema(schema, text, graph): source_label = edge["source_vertex_label"] source_id = f"{source_label}-{s}" if source_id not in vertices_dict: - vertices_dict[source_id] = {"id": source_id, "name": s, "label": source_label, - "properties": {}} + vertices_dict[source_id] = { + "id": source_id, + "name": s, + "label": source_label, + "properties": {}, + } target_label = edge["target_vertex_label"] target_id = f"{target_label}-{o}" if target_id not in vertices_dict: - vertices_dict[target_id] = {"id": target_id, "name": o, "label": target_label, - "properties": {}} - graph["edges"].append({"start": source_id, "end": target_id, "type": label, - "properties": {}}) + vertices_dict[target_id] = { + "id": target_id, + "name": o, + "label": target_label, + "properties": {}, + } + graph["edges"].append( + { + "start": source_id, + "end": target_id, + "type": label, + "properties": {}, + } + ) break - graph["vertices"] = vertices_dict.values() + graph["vertices"] = list(vertices_dict.values()) class InfoExtract: - def __init__( - self, - llm: BaseLLM, - example_prompt: Optional[str] = None - ) -> None: + def __init__(self, llm: BaseLLM, example_prompt: Optional[str] = None) -> None: self.llm = llm self.example_prompt = example_prompt @@ -152,7 +176,12 @@ def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]: for sentence in chunks: proceeded_chunk = self.extract_triples_by_llm(schema, sentence) - log.debug("[Legacy] %s input: %s \n output:%s", self.__class__.__name__, sentence, proceeded_chunk) + log.debug( + "[Legacy] %s input: %s \n output:%s", + self.__class__.__name__, + sentence, + proceeded_chunk, + ) if schema: extract_triples_by_regex_with_schema(schema, proceeded_chunk, context) else: @@ -175,7 +204,152 @@ def valid(self, element_id: str, max_length: int = 256) -> bool: return True def _filter_long_id(self, graph) -> Dict[str, List[Any]]: - graph["vertices"] = [vertex for vertex in graph["vertices"] if self.valid(vertex["id"])] - graph["edges"] = [edge for edge in graph["edges"] - if self.valid(edge["start"]) and self.valid(edge["end"])] + graph["vertices"] = [ + vertex for vertex in graph["vertices"] if self.valid(vertex["id"]) + ] + graph["edges"] = [ + edge + for edge in graph["edges"] + if self.valid(edge["start"]) and self.valid(edge["end"]) + ] return graph + + +class InfoExtractNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + return init_context(self) + + def node_init(self): + self.llm = get_chat_llm(llm_settings) + if self.wk_input.example_prompt is None: + return CStatus(-1, "Error occurs when prepare for workflow input") + self.example_prompt = self.wk_input.example_prompt + return CStatus() + + def extract_triples_by_regex_with_schema(self, schema, text): + text = text.replace("\\n", " ").replace("\\", " ").replace("\n", " ") + pattern = r"\((.*?), (.*?), (.*?)\) - ([^ ]*)" + matches = re.findall(pattern, text) + + vertices_dict = {v["id"]: v for v in self.context.vertices} + for match in matches: + s, p, o, label = [item.strip() for item in match] + if None in [label, s, p, o]: + continue + # TODO: use a more efficient way to compare the extract & input property + p_lower = p.lower() + for vertex in schema["vertices"]: + if vertex["vertex_label"] == label and any( + pp.lower() == p_lower for pp in vertex["properties"] + ): + id = f"{label}-{s}" + if id not in vertices_dict: + vertices_dict[id] = { + "id": id, + "name": s, + "label": label, + "properties": {p: o}, + } + else: + vertices_dict[id]["properties"].update({p: o}) + break + for edge in schema["edges"]: + if edge["edge_label"] == label: + source_label = edge["source_vertex_label"] + source_id = f"{source_label}-{s}" + if source_id not in vertices_dict: + vertices_dict[source_id] = { + "id": source_id, + "name": s, + "label": source_label, + "properties": {}, + } + target_label = edge["target_vertex_label"] + target_id = f"{target_label}-{o}" + if target_id not in vertices_dict: + vertices_dict[target_id] = { + "id": target_id, + "name": o, + "label": target_label, + "properties": {}, + } + self.context.edges.append( + { + "start": source_id, + "end": target_id, + "type": label, + "properties": {}, + } + ) + break + self.context.vertices = list(vertices_dict.values()) + + def extract_triples_by_regex(self, text): + text = text.replace("\\n", " ").replace("\\", " ").replace("\n", " ") + pattern = r"\((.*?), (.*?), (.*?)\)" + self.context.triples += re.findall(pattern, text) + + def run(self) -> CStatus: + sts = self.node_init() + if sts.isErr(): + return sts + self.context.lock() + if self.context.chunks is None: + self.context.unlock() + raise ValueError("parameter required by extract node not found in context.") + schema = self.context.schema + chunks = self.context.chunks + + if schema: + self.context.vertices = [] + self.context.edges = [] + else: + self.context.triples = [] + + self.context.unlock() + + for sentence in chunks: + proceeded_chunk = self.extract_triples_by_llm(schema, sentence) + log.debug( + "[Legacy] %s input: %s \n output:%s", + self.__class__.__name__, + sentence, + proceeded_chunk, + ) + if schema: + self.extract_triples_by_regex_with_schema(schema, proceeded_chunk) + else: + self.extract_triples_by_regex(proceeded_chunk) + + if self.context.call_count: + self.context.call_count += len(chunks) + else: + self.context.call_count = len(chunks) + self._filter_long_id() + return CStatus() + + def extract_triples_by_llm(self, schema, chunk) -> str: + prompt = generate_extract_triple_prompt(chunk, schema) + if self.example_prompt is not None: + prompt = self.example_prompt + prompt + return self.llm.generate(prompt=prompt) + + # TODO: make 'max_length' be a configurable param in settings.py/settings.cfg + def valid(self, element_id: str, max_length: int = 256) -> bool: + if len(element_id.encode("utf-8")) >= max_length: + log.warning("Filter out GraphElementID too long: %s", element_id) + return False + return True + + def _filter_long_id(self): + self.context.vertices = [ + vertex for vertex in self.context.vertices if self.valid(vertex["id"]) + ] + self.context.edges = [ + edge + for edge in self.context.edges + if self.valid(edge["start"]) and self.valid(edge["end"]) + ] 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 faff1c6b2..6e492b8f5 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 @@ -21,16 +21,19 @@ import re from typing import List, Any, Dict -from hugegraph_llm.config import prompt +from hugegraph_llm.config import llm_settings, prompt from hugegraph_llm.document.chunk_split import ChunkSplitter from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.utils.log import log -""" -TODO: It is not clear whether there is any other dependence on the SCHEMA_EXAMPLE_PROMPT variable. -Because the SCHEMA_EXAMPLE_PROMPT variable will no longer change based on -prompt.extract_graph_prompt changes after the system loads, this does not seem to meet expectations. -""" +from hugegraph_llm.operators.util import init_context +from hugegraph_llm.models.llms.init_llm import get_chat_llm +from hugegraph_llm.state.ai_state import WkFlowState, WkFlowInput +from PyCGraph import GNode, CStatus + +# TODO: It is not clear whether there is any other dependence on the SCHEMA_EXAMPLE_PROMPT variable. +# Because the SCHEMA_EXAMPLE_PROMPT variable will no longer change based on +# prompt.extract_graph_prompt changes after the system loads, this does not seem to meet expectations. SCHEMA_EXAMPLE_PROMPT = prompt.extract_graph_prompt @@ -60,20 +63,18 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: properties_map["vertex"][vertex["name"]] = { "primary_keys": vertex["primary_keys"], "nullable_keys": vertex["nullable_keys"], - "properties": vertex["properties"] + "properties": vertex["properties"], } for edge in schema["edgelabels"]: - properties_map["edge"][edge["name"]] = { - "properties": edge["properties"] - } + properties_map["edge"][edge["name"]] = {"properties": edge["properties"]} log.info("properties_map: %s", properties_map) for item in items: item_type = item["type"] if item_type == "vertex": label = item["label"] - non_nullable_keys = ( - set(properties_map[item_type][label]["properties"]) - .difference(set(properties_map[item_type][label]["nullable_keys"]))) + non_nullable_keys = set( + properties_map[item_type][label]["properties"] + ).difference(set(properties_map[item_type][label]["nullable_keys"])) for key in non_nullable_keys: if key not in item["properties"]: item["properties"][key] = "NULL" @@ -87,9 +88,7 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: class PropertyGraphExtract: def __init__( - self, - llm: BaseLLM, - example_prompt: str = prompt.extract_graph_prompt + self, llm: BaseLLM, example_prompt: str = prompt.extract_graph_prompt ) -> None: self.llm = llm self.example_prompt = example_prompt @@ -105,7 +104,12 @@ def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]: items = [] for chunk in chunks: proceeded_chunk = self.extract_property_graph_by_llm(schema, chunk) - log.debug("[LLM] %s input: %s \n output:%s", self.__class__.__name__, chunk, proceeded_chunk) + log.debug( + "[LLM] %s input: %s \n output:%s", + self.__class__.__name__, + chunk, + proceeded_chunk, + ) items.extend(self._extract_and_filter_label(schema, proceeded_chunk)) items = filter_item(schema, items) for item in items: @@ -125,10 +129,132 @@ def extract_property_graph_by_llm(self, schema, chunk): def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: # Use regex to extract a JSON object with curly braces - json_match = re.search(r'({.*})', text, re.DOTALL) + json_match = re.search(r"({.*})", text, re.DOTALL) + if not json_match: + log.critical( + "Invalid property graph! No JSON object found, " + "please check the output format example in prompt." + ) + return [] + json_str = json_match.group(1).strip() + + items = [] + try: + property_graph = json.loads(json_str) + # 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'." + ) + return items + + # Create sets for valid vertex and edge labels based on the schema + vertex_label_set = {vertex["name"] for vertex in schema["vertexlabels"]} + edge_label_set = {edge["name"] for edge in schema["edgelabels"]} + + def process_items(item_list, valid_labels, item_type): + for item in item_list: + if not isinstance(item, dict): + log.warning( + "Invalid property graph item type '%s'.", type(item) + ) + continue + if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): + log.warning("Invalid item keys '%s'.", item.keys()) + continue + if item["label"] not in valid_labels: + log.warning( + "Invalid %s label '%s' has been ignored.", + item_type, + item["label"], + ) + continue + items.append(item) + + process_items(property_graph["vertices"], vertex_label_set, "vertex") + process_items(property_graph["edges"], edge_label_set, "edge") + except json.JSONDecodeError: + log.critical( + "Invalid property graph JSON! Please check the extracted JSON data carefully" + ) + return items + + +class PropertyGraphExtractNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name + return init_context(self) + + def node_init(self): + self.llm = get_chat_llm(llm_settings) + if self.wk_input.example_prompt is None: + return CStatus(-1, "Error occurs when prepare for workflow input") + self.example_prompt = self.wk_input.example_prompt + return CStatus() + + def run(self) -> CStatus: + sts = self.node_init() + if sts.isErr(): + return sts + self.context.lock() + try: + if self.context.schema is None or self.context.chunks is None: + raise ValueError( + "parameter required by extract node not found in context." + ) + schema = self.context.schema + chunks = self.context.chunks + if self.context.vertices is None: + self.context.vertices = [] + if self.context.edges is None: + self.context.edges = [] + finally: + self.context.unlock() + + items = [] + for chunk in chunks: + proceeded_chunk = self.extract_property_graph_by_llm(schema, chunk) + log.debug( + "[LLM] %s input: %s \n output:%s", + self.__class__.__name__, + chunk, + proceeded_chunk, + ) + items.extend(self._extract_and_filter_label(schema, proceeded_chunk)) + items = filter_item(schema, items) + self.context.lock() + try: + for item in items: + if item["type"] == "vertex": + self.context.vertices.append(item) + elif item["type"] == "edge": + self.context.edges.append(item) + finally: + self.context.unlock() + self.context.call_count = (self.context.call_count or 0) + len(chunks) + return CStatus() + + def extract_property_graph_by_llm(self, schema, chunk): + prompt = generate_extract_property_graph_prompt(chunk, schema) + if self.example_prompt is not None: + prompt = self.example_prompt + prompt + return self.llm.generate(prompt=prompt) + + def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: + # Use regex to extract a JSON object with curly braces + json_match = re.search(r"({.*})", text, re.DOTALL) if not json_match: - log.critical("Invalid property graph! No JSON object found, " - "please check the output format example in prompt.") + log.critical( + "Invalid property graph! No JSON object found, " + "please check the output format example in prompt." + ) return [] json_str = json_match.group(1).strip() @@ -136,8 +262,14 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: try: property_graph = json.loads(json_str) # 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 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'." + ) return items # Create sets for valid vertex and edge labels based on the schema @@ -147,18 +279,26 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: def process_items(item_list, valid_labels, item_type): for item in item_list: if not isinstance(item, dict): - log.warning("Invalid property graph item type '%s'.", type(item)) + log.warning( + "Invalid property graph item type '%s'.", type(item) + ) continue if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): log.warning("Invalid item keys '%s'.", item.keys()) continue if item["label"] not in valid_labels: - log.warning("Invalid %s label '%s' has been ignored.", item_type, item["label"]) + log.warning( + "Invalid %s label '%s' has been ignored.", + item_type, + item["label"], + ) continue items.append(item) process_items(property_graph["vertices"], vertex_label_set, "vertex") process_items(property_graph["edges"], edge_label_set, "edge") except json.JSONDecodeError: - log.critical("Invalid property graph JSON! Please check the extracted JSON data carefully") + log.critical( + "Invalid property graph JSON! Please check the extracted JSON data carefully" + ) return items diff --git a/hugegraph-llm/src/hugegraph_llm/operators/util.py b/hugegraph-llm/src/hugegraph_llm/operators/util.py new file mode 100644 index 000000000..60bdc2e86 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/operators/util.py @@ -0,0 +1,27 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from PyCGraph import CStatus + + +def init_context(obj) -> CStatus: + try: + obj.context = obj.getGParamWithNoEmpty("wkflow_state") + obj.wk_input = obj.getGParamWithNoEmpty("wkflow_input") + if obj.context is None or obj.wk_input is None: + return CStatus(-1, "Required workflow parameters not found") + return CStatus() + except Exception as e: + return CStatus(-1, f"Failed to initialize context: {str(e)}") diff --git a/hugegraph-llm/src/hugegraph_llm/state/__init__.py b/hugegraph-llm/src/hugegraph_llm/state/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/state/__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/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py new file mode 100644 index 000000000..0543aa2b4 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -0,0 +1,81 @@ +# 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 PyCGraph import GParam, CStatus + +from typing import Union, List, Optional, Any + + +class WkFlowInput(GParam): + texts: Union[str, List[str]] = None # texts input used by ChunkSplit Node + language: str = None # language configuration used by ChunkSplit Node + split_type: str = None # split type used by ChunkSplit Node + example_prompt: str = None # need by graph information extract + schema: str = None # Schema information requeired by SchemaNode + graph_name: str = None + + def reset(self, _: CStatus) -> None: + self.texts = None + self.language = None + self.split_type = None + self.example_prompt = None + self.schema = None + self.graph_name = None + + +class WkFlowState(GParam): + schema: Optional[str] = None # schema message + simple_schema: Optional[str] = None + chunks: Optional[List[str]] = None + edges: Optional[List[Any]] = None + vertices: Optional[List[Any]] = None + triples: Optional[List[Any]] = None + call_count: Optional[int] = None + + keywords: Optional[List[str]] = None + vector_result = None + graph_result = None + keywords_embeddings = None + + def setup(self): + self.schema = None + self.simple_schema = None + self.chunks = None + self.edges = None + self.vertices = None + self.triples = None + self.call_count = None + + self.keywords = None + self.vector_result = None + self.graph_result = None + self.keywords_embeddings = None + + return CStatus() + + def to_json(self): + """ + Automatically returns a JSON-formatted dictionary of all non-None instance members, + eliminating the need to manually maintain the member list. + + Returns: + dict: A dictionary containing non-None instance members and their serialized values. + """ + # Only export instance attributes (excluding methods and class attributes) whose values are not None + return { + k: v + for k, v in self.__dict__.items() + if not k.startswith("_") and v is not None + } diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 9fef06d2b..f61b5f843 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -22,6 +22,7 @@ from typing import Dict, Any, Union, Optional import gradio as gr +from hugegraph_llm.flows.scheduler import SchedulerSingleton from .embedding_utils import get_filename_prefix, get_index_folder_name from .hugegraph_utils import get_hg_client, clean_hg_data @@ -35,11 +36,17 @@ def get_graph_index_info(): - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) graph_summary_info = builder.fetch_graph_data().run() - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) - filename_prefix = get_filename_prefix(llm_settings.embedding_type, getattr(builder.embedding, "model_name", None)) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(builder.embedding, "model_name", None) + ) vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) graph_summary_info["vid_index"] = { "embed_dim": vector_index.index.d, @@ -50,15 +57,20 @@ def get_graph_index_info(): def clean_all_graph_index(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix(llm_settings.embedding_type, - getattr(Embeddings().get_embedding(), "model_name", None)) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, + getattr(Embeddings().get_embedding(), "model_name", None), + ) VectorIndex.clean( - str(os.path.join(resource_path, folder_name, "graph_vids")), - filename_prefix) + str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix + ) VectorIndex.clean( str(os.path.join(resource_path, folder_name, "gremlin_examples")), - filename_prefix) + filename_prefix, + ) log.warning("Clear graph index and text2gql index successfully!") gr.Info("Clear graph index and text2gql index successfully!") @@ -71,7 +83,7 @@ def clean_all_graph_data(): def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: schema = schema.strip() - if schema.startswith('{'): + if schema.startswith("{"): try: schema = json.loads(schema) builder.import_schema(from_user_defined=schema) @@ -84,16 +96,20 @@ def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: return None -def extract_graph(input_file, input_text, schema, example_prompt) -> str: +def extract_graph_origin(input_file, input_text, schema, example_prompt) -> str: texts = read_documents(input_file, input_text) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) if not schema: return "ERROR: please input with correct schema/format." error_message = parse_schema(schema, builder) if error_message: return error_message - builder.chunk_split(texts, "document", "zh").extract_info(example_prompt, "property_graph") + builder.chunk_split(texts, "document", "zh").extract_info( + example_prompt, "property_graph" + ) try: context = builder.run() @@ -103,19 +119,40 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: { "vertices": context["vertices"], "edges": context["edges"], - "warning": "The schema may not match the Doc" + "warning": "The schema may not match the Doc", }, ensure_ascii=False, - indent=2 + indent=2, ) - return json.dumps({"vertices": context["vertices"], "edges": context["edges"]}, ensure_ascii=False, indent=2) + return json.dumps( + {"vertices": context["vertices"], "edges": context["edges"]}, + ensure_ascii=False, + indent=2, + ) + except Exception as e: # pylint: disable=broad-exception-caught + log.error(e) + raise gr.Error(str(e)) + + +def extract_graph(input_file, input_text, schema, example_prompt) -> str: + texts = read_documents(input_file, input_text) + scheduler = SchedulerSingleton.get_instance() + if not schema: + return "ERROR: please input with correct schema/format." + + try: + return scheduler.schedule_flow( + "graph_extract", schema, texts, example_prompt, "property_graph" + ) except Exception as e: # pylint: disable=broad-exception-caught log.error(e) raise gr.Error(str(e)) def update_vid_embedding(): - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) builder.fetch_graph_data().build_vertex_id_semantic_index() log.debug("Operators: %s", builder.operators) try: @@ -132,7 +169,9 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: data_json = json.loads(data.strip()) log.debug("Import graph data: %s", data) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) if schema: error_message = parse_schema(schema, builder) if error_message: @@ -154,7 +193,7 @@ def build_schema(input_text, query_example, few_shot): context = { "raw_texts": [input_text] if input_text else [], "query_examples": [], - "few_shot_schema": {} + "few_shot_schema": {}, } if few_shot: @@ -170,7 +209,7 @@ def build_schema(input_text, query_example, few_shot): context["query_examples"] = [ { "description": ex.get("description", ""), - "gremlin": ex.get("gremlin", "") + "gremlin": ex.get("gremlin", ""), } for ex in parsed_examples if isinstance(ex, dict) and "description" in ex and "gremlin" in ex @@ -178,7 +217,9 @@ def build_schema(input_text, query_example, few_shot): except json.JSONDecodeError as e: raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) try: schema = builder.build_schema().run(context) except Exception as e: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 62bcdd9cb..138b0d359 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -23,11 +23,12 @@ from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index import VectorIndex -from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.kg_construction_task import KgBuilder -from hugegraph_llm.utils.embedding_utils import get_filename_prefix, get_index_folder_name -from hugegraph_llm.utils.hugegraph_utils import get_hg_client +from hugegraph_llm.models.embeddings.init_embedding import model_map +from hugegraph_llm.flows.scheduler import SchedulerSingleton +from hugegraph_llm.utils.embedding_utils import ( + get_filename_prefix, + get_index_folder_name, +) def read_documents(input_file, input_text): @@ -49,7 +50,9 @@ def read_documents(input_file, input_text): texts.append(text) elif full_path.endswith(".pdf"): # TODO: support PDF file - raise gr.Error("PDF will be supported later! Try to upload text/docx now") + raise gr.Error( + "PDF will be supported later! Try to upload text/docx now" + ) else: raise gr.Error("Please input txt or docx file.") else: @@ -59,33 +62,44 @@ def read_documents(input_file, input_text): # pylint: disable=C0301 def get_vector_index_info(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix(llm_settings.embedding_type, - getattr(Embeddings().get_embedding(), "model_name", None)) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) + ) chunk_vector_index = VectorIndex.from_index_file( str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix, - record_miss=False + record_miss=False, ) graph_vid_vector_index = VectorIndex.from_index_file( - str(os.path.join(resource_path, folder_name, "graph_vids")), - filename_prefix + str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix + ) + return json.dumps( + { + "embed_dim": chunk_vector_index.index.d, + "vector_info": { + "chunk_vector_num": chunk_vector_index.index.ntotal, + "graph_vid_vector_num": graph_vid_vector_index.index.ntotal, + "graph_properties_vector_num": len(chunk_vector_index.properties), + }, + }, + ensure_ascii=False, + indent=2, ) - return json.dumps({ - "embed_dim": chunk_vector_index.index.d, - "vector_info": { - "chunk_vector_num": chunk_vector_index.index.ntotal, - "graph_vid_vector_num": graph_vid_vector_index.index.ntotal, - "graph_properties_vector_num": len(chunk_vector_index.properties) - } - }, ensure_ascii=False, indent=2) def clean_vector_index(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix(llm_settings.embedding_type, - getattr(Embeddings().get_embedding(), "model_name", None)) - VectorIndex.clean(str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) + ) + VectorIndex.clean( + str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix + ) gr.Info("Clean vector index successfully!") @@ -93,6 +107,5 @@ def build_vector_index(input_file, input_text): if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - context = builder.chunk_split(texts, "paragraph", "zh").build_vector_index().run() - return json.dumps(context, ensure_ascii=False, indent=2) + scheduler = SchedulerSingleton.get_instance() + return scheduler.schedule_flow("build_vector_index", texts) From 41aeae51f83965fcc36ac4e9589cf5a5f898a33a Mon Sep 17 00:00:00 2001 From: Linyu <94553312+weijinglin@users.noreply.github.com> Date: Thu, 25 Sep 2025 15:09:38 +0800 Subject: [PATCH 02/71] Refactor: Refactor hugegraph-ai to using CGraph & port some usecases in web demo (#49) --- .../spec/hugegraph-llm/fixed_flow/design.md | 643 ++++++++++++++++++ .../hugegraph-llm/fixed_flow/requirements.md | 24 + .../spec/hugegraph-llm/fixed_flow/tasks.md | 36 + .../demo/rag_demo/vector_graph_block.py | 21 +- .../src/hugegraph_llm/flows/build_schema.py | 71 ++ .../hugegraph_llm/flows/build_vector_index.py | 4 +- .../flows/get_graph_index_info.py | 68 ++ .../src/hugegraph_llm/flows/graph_extract.py | 58 +- .../hugegraph_llm/flows/import_graph_data.py | 65 ++ .../hugegraph_llm/flows/prompt_generate.py | 63 ++ .../src/hugegraph_llm/flows/scheduler.py | 31 +- .../flows/update_vid_embeddings.py | 47 ++ .../src/hugegraph_llm/flows/utils.py | 34 + .../src/hugegraph_llm/nodes/base_node.py | 71 ++ .../nodes/document_node/chunk_split.py | 43 ++ .../hugegraph_node/commit_to_hugegraph.py | 35 + .../nodes/hugegraph_node/fetch_graph_data.py | 33 + .../nodes/hugegraph_node/schema.py | 74 ++ .../nodes/index_node/build_semantic_index.py | 34 + .../nodes/index_node/build_vector_index.py | 34 + .../nodes/llm_node/extract_info.py | 52 ++ .../nodes/llm_node/prompt_generate.py | 59 ++ .../nodes/llm_node/schema_build.py | 91 +++ hugegraph-llm/src/hugegraph_llm/nodes/util.py | 27 + .../operators/common_op/check_schema.py | 160 ----- .../operators/document_op/chunk_split.py | 58 -- .../hugegraph_op/commit_to_hugegraph.py | 127 +++- .../operators/hugegraph_op/schema_manager.py | 75 -- .../operators/index_op/build_vector_index.py | 48 -- .../operators/llm_op/info_extract.py | 146 ---- .../llm_op/property_graph_extract.py | 127 +--- .../src/hugegraph_llm/state/ai_state.py | 28 + .../hugegraph_llm/utils/graph_index_utils.py | 40 ++ 33 files changed, 1811 insertions(+), 716 deletions(-) create mode 100644 .vibedev/spec/hugegraph-llm/fixed_flow/design.md create mode 100644 .vibedev/spec/hugegraph-llm/fixed_flow/requirements.md create mode 100644 .vibedev/spec/hugegraph-llm/fixed_flow/tasks.md create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/build_schema.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/utils.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/base_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/document_node/chunk_split.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/commit_to_hugegraph.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/fetch_graph_data.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_semantic_index.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/extract_info.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/util.py diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/design.md b/.vibedev/spec/hugegraph-llm/fixed_flow/design.md new file mode 100644 index 000000000..c5777236d --- /dev/null +++ b/.vibedev/spec/hugegraph-llm/fixed_flow/design.md @@ -0,0 +1,643 @@ +# Hugegraph-ai 固定工作流执行引擎设计文档 + +## 概述 + +Hugegraph固定工作流执行引擎是用来执行固定工作流的工作流执行引擎，每个工作流对应到实际Web Demo的一个具体用例，包括向量索引的构建，图索引的构建等等。该引擎基于PyCGraph框架构建，提供了高性能、可复用的流水线调度能力。 + +### 设计目标 + +- **性能优异**：通过流水线复用机制保证固定工作流的执行性能 +- **高可靠性**：确保数据一致性和故障恢复能力，提供完善的错误处理机制 +- **易于扩展**：能够简单轻松地新增固定工作流，支持动态调度 +- **资源优化**：通过流水线池化管理，减少重复构图开销 + +### 技术栈 + +- **PyCGraph**：基于C++的高性能图计算框架，提供GPipeline和GPipelineManager +- **Python**：主要开发语言，提供业务逻辑和接口层 +- **Threading**：支持并发调度和线程安全 + +### 模块分层 +```text +hugegraph-llm/ +└── src/ + └── hugegraph_llm/ + ├── api/ # FastAPI 接口层，提供 rag_api、admin_api 等服务 + ├── config/ # 配置管理，包含各类配置与生成工具 + ├── demo/ # Gradio Web Demo 及相关交互应用 + ├── document/ # 文档处理与分块等工具 + ├── enums/ # 枚举类型定义 + ├── flows/ # 工作流调度与核心流程（如向量/图索引构建、数据导入等） + │ ├── __init__.py + │ ├── common.py # BaseFlow抽象基类 + │ ├── scheduler.py # 调度器核心实现 + │ ├── build_vector_index.py # 向量索引构建工作流 + │ ├── graph_extract.py # 图抽取工作流 + │ ├── import_graph_data.py # 图数据导入工作流 + │ ├── update_vid_embeddings.py # 向量更新工作流 + │ ├── get_graph_index_info.py # 图索引信息获取工作流 + │ ├── build_schema.py # 模式构建工作流 + │ └── prompt_generate.py # 提示词生成工作流 + ├── indices/ # 各类索引实现（向量、图、关键词等） + ├── middleware/ # 中间件与请求处理 + ├── models/ # LLM、Embedding、Reranker 等模型相关 + ├── nodes/ # Node调度层，负责Operator生命周期和上下文管理 + │ ├── base_node.py + │ ├── document_node/ + │ ├── hugegraph_node/ + │ ├── index_node/ + │ ├── llm_node/ + │ └── util.py + ├── operators/ # 主要算子与任务（如 KG 构建、GraphRAG、Text2Gremlin 等） + ├── resources/ # 资源文件（Prompt、示例、Gremlin 模板等） + ├── state/ # 状态管理 + ├── utils/ # 工具类与通用方法 + └── __init__.py # 包初始化 +``` + +## 架构设计 + +### 整体架构 + +> 新架构在Flow与Operator之间引入Node层，Node负责Operator的生命周期管理、上下文绑定、参数区解耦和并发安全，所有Flow均通过Node组装，Operator只关注业务实现。 + +#### 架构图 + +```mermaid +graph TB + subgraph UserLayer["用户层"] + User["用户请求"] + end + + subgraph SchedulerLayer["调度层"] + Scheduler["Scheduler
调度器"] + Singleton["SchedulerSingleton
单例管理器"] + end + + subgraph FlowLayer["工作流层"] + Pool["pipeline_pool
流水线池"] + BVI["BuildVectorIndexFlow
向量索引构建"] + GE["GraphExtractFlow
图抽取工作流"] + end + + subgraph PyCGraphLayer["PyCGraph层"] + Manager1["GPipelineManager
向量索引管理器"] + Manager2["GPipelineManager
图抽取管理器"] + Pipeline1["GPipeline
向量索引流水线"] + Pipeline2["GPipeline
图抽取流水线"] + end + + subgraph OperatorLayer["算子层"] + ChunkSplit["ChunkSplitNode
文档分块"] + BuildVector["BuildVectorIndexNode
向量索引构建"] + SchemaNode["SchemaNode
模式管理"] + InfoExtract["ExtractNode
信息抽取"] + PropGraph["Commit2GraphNode
图数据导入"] + FetchNode["FetchGraphDataNode
图数据拉取"] + SemanticIndex["BuildSemanticIndexNode
语义索引构建"] + end + + subgraph StateLayer["状态层"] + WkInput["wkflow_input
工作流输入"] + WkState["wkflow_state
工作流状态"] + end + + User --> Scheduler + Scheduler --> Singleton + Scheduler --> Pool + Pool --> BVI + Pool --> GE + BVI --> Manager1 + GE --> Manager2 + Manager1 --> Pipeline1 + Manager2 --> Pipeline2 + Pipeline1 --> ChunkSplit + Pipeline1 --> BuildVector + Pipeline2 --> SchemaNode + Pipeline2 --> ChunkSplit + Pipeline2 --> InfoExtract + Pipeline2 --> PropGraph + Pipeline1 --> WkInput + Pipeline1 --> WkState + Pipeline2 --> WkInput + Pipeline2 --> WkState + + style Scheduler fill:#e1f5fe + style Pool fill:#f3e5f5 + style Manager1 fill:#fff3e0 + style Manager2 fill:#fff3e0 + style Pipeline1 fill:#e8f5e8 + style Pipeline2 fill:#e8f5e8 +``` + +#### 调度流程图 + +```mermaid +flowchart TD + Start([开始]) --> CheckFlow{检查工作流
是否支持} + CheckFlow -->|否| Error1[抛出ValueError] + CheckFlow -->|是| FetchPipeline[从Manager获取
可复用Pipeline] + + FetchPipeline --> IsNull{Pipeline
是否为null} + + IsNull -->|是| BuildNew[构建新Pipeline] + BuildNew --> InitPipeline[初始化Pipeline] + InitPipeline --> InitCheck{初始化
是否成功} + InitCheck -->|否| Error2[记录错误并中止] + InitCheck -->|是| RunPipeline[执行Pipeline] + RunPipeline --> RunCheck{执行
是否成功} + RunCheck -->|否| Error3[记录错误并中止] + RunCheck -->|是| PostDeal[后处理结果] + PostDeal --> AddToPool[添加到复用池] + AddToPool --> Return[返回结果] + + IsNull -->|否| PrepareInput[准备输入数据] + PrepareInput --> RunReused[执行复用Pipeline] + RunReused --> ReusedCheck{执行
是否成功} + ReusedCheck -->|否| Error4[抛出RuntimeError] + ReusedCheck -->|是| PostDealReused[后处理结果] + PostDealReused --> ReleasePipeline[释放Pipeline] + ReleasePipeline --> Return + + Error1 --> End([结束]) + Error2 --> End + Error3 --> End + Error4 --> End + Return --> End + + style Start fill:#4caf50 + style End fill:#f44336 + style CheckFlow fill:#ff9800 + style IsNull fill:#ff9800 + style InitCheck fill:#ff9800 + style RunCheck fill:#ff9800 + style ReusedCheck fill:#ff9800 +``` + +### 核心组件 + +#### 1. Scheduler（调度器） +- **职责**：调度中心，维护 `pipeline_pool`，提供统一的工作流调度接口 +- **特性**： + - 支持多种工作流类型（build_vector_index、graph_extract、import_graph_data、update_vid_embeddings、get_graph_index_info、build_schema、prompt_generate等） + - 流水线池化管理，支持复用 + - 线程安全的单例模式 + - 可配置的最大流水线数量 + +#### 2. GPipelineManager（流水线管理器） +- **来源**：PyCGraph框架提供 +- **职责**：负责流水线对象 `GPipeline` 的获取、添加、释放与复用 +- **特性**： + - 自动管理流水线生命周期 + - 支持流水线复用和资源回收 + - 提供fetch/add/release操作接口 + +#### 3. BaseFlow（工作流基类） +- **职责**：工作流构建与前后处理抽象 +- **接口**： + - `prepare()`: 预处理接口，准备输入数据 + - `build_flow()`: 组装Node并注册依赖关系 + - `post_deal()`: 后处理接口，处理执行结果 +- **实现**： + - `BuildVectorIndexFlow`: 向量索引构建工作流 + - `GraphExtractFlow`: 图抽取工作流 + - `ImportGraphDataFlow`: 图数据导入工作流 + - `UpdateVidEmbeddingsFlows`: 向量更新工作流 + - `GetGraphIndexInfoFlow`: 图索引信息获取工作流 + - `BuildSchemaFlow`: 模式构建工作流 + - `PromptGenerateFlow`: 提示词生成工作流 + +#### 4. Node（节点调度器） +- **职责**：作为Operator的生命周期管理者，负责参数区绑定、上下文初始化、并发安全、异常处理等。 +- **特性**： + - 统一生命周期接口（init、node_init、run、operator_schedule） + - 通过参数区（wkflow_input/wkflow_state）与Flow/Operator解耦 + - Operator只需实现run(data_json)方法，Node负责调度和结果写回 + - 典型Node如：ChunkSplitNode、BuildVectorIndexNode、SchemaNode、ExtractNode、Commit2GraphNode、FetchGraphDataNode、BuildSemanticIndexNode、SchemaBuildNode、PromptGenerateNode等 + +#### 5. Operator（算子） +- **职责**：实现具体的业务原子操作 +- **特性**： + - 只需关注自身业务逻辑实现 + - 由Node统一调度 + +#### 6. GPipeline（流水线实例） +- **来源**：PyCGraph框架提供 +- **职责**：具体流水线实例，包含参数区与节点DAG拓扑 +- **参数区**： + - `wkflow_input`: 流水线运行输入 + - `wkflow_state`: 流水线运行状态与中间结果 + +### 核心数据结构 + +```python +# Scheduler核心数据结构 +Scheduler.pipeline_pool: Dict[str, Any] = { + "build_vector_index": { + "manager": GPipelineManager(), + "flow": BuildVectorIndexFlow(), + }, + "graph_extract": { + "manager": GPipelineManager(), + "flow": GraphExtractFlow(), + } +} +``` + +### 调度流程 + +#### schedule_flow方法执行流程 + +1. **工作流验证**：校验 `flow` 是否受支持，查表获取对应的 `manager` 与 `flow` 实例 + +2. **流水线获取**：从 `manager.fetch()` 获取可复用的 `GPipeline` + +3. **新流水线处理**（当fetch()返回None时）： + - 调用 `flow.build_flow(*args, **kwargs)` 构建新流水线 + - 调用 `pipeline.init()` 完成初始化，失败则记录错误并中止 + - 调用 `pipeline.run()` 执行，失败则中止 + - 调用 `flow.post_deal(pipeline)` 生成输出 + - 调用 `manager.add(pipeline)` 将流水线加入可复用池 + +4. **复用流水线处理**（当fetch()返回现有流水线时）： + - 从 `pipeline.getGParamWithNoEmpty("wkflow_input")` 获取输入对象 + - 调用 `flow.prepare(prepared_input, *args, **kwargs)` 进行参数刷新 + - 调用 `pipeline.run()` 执行，失败则中止 + - 调用 `flow.post_deal(pipeline)` 生成输出 + - 调用 `manager.release(pipeline)` 归还流水线 + +### 并发与复用策略 + +#### 线程安全 +- `SchedulerSingleton` 使用双重检查锁保证全局单例 +- 线程安全获取 `Scheduler` 实例 + +#### 资源管理 +- 每种 `flow` 拥有独立的 `GPipelineManager` +- 最大并发量由 `Scheduler.max_pipeline` 与底层 `GPipelineManager` 策略共同约束 +- 通过 `fetch/add/release` 机制减少重复构图的开销 + +#### 性能优化 +- 流水线复用机制适合高频相同工作流场景 +- 减少重复初始化和构图的时间开销 +- 支持并发执行多个工作流实例 + +### 错误处理与日志 + +#### 错误检测 +- 对 `init/run` 的 `Status.isErr()` 进行检测 +- 统一抛出 `RuntimeError` 并记录详细 `status.getInfo()` +- 提供完整的错误堆栈信息 + +#### 日志记录 +- 使用统一的日志系统记录关键操作 +- 记录流水线执行状态和错误信息 +- 支持不同级别的日志输出 + +#### 结果处理 +- `flow.post_deal` 负责将 `wkflow_state` 转换为对外可消费结果（如JSON） +- 提供标准化的输出格式 +- 支持错误信息的友好展示 + +### 扩展指引 + +#### 新增Node/Operator/Flow步骤 +1. 实现Operator业务逻辑（如ChunkSplit/BuildVectorIndex/InfoExtract等） +2. 实现对应Node（继承BaseNode，负责参数区绑定和调度Operator） +3. 在Flow中组装Node，注册依赖关系 +4. 在Scheduler注册新的Flow + +#### 输入输出约定 +- 统一使用 `wkflow_input` 作为输入载体 +- 统一使用 `wkflow_state` 作为状态与结果容器 +- 确保可复用流水线在不同请求间可被快速重置 + +#### 最佳实践 +- 保持Flow类的无状态设计 +- 合理使用流水线复用机制 +- 提供完善的错误处理和日志记录 +- 遵循统一的接口规范 + +## Flow对象设计 + +### BaseFlow抽象基类 + +```python +class BaseFlow(ABC): + """ + Base class for flows, defines three interface methods: prepare, build_flow, and post_deal. + """ + + @abstractmethod + def prepare(self, prepared_input: WkFlowInput, *args, **kwargs): + """ + Pre-processing interface. + """ + pass + + @abstractmethod + def build_flow(self, *args, **kwargs): + """ + Interface for building the flow. + """ + pass + + @abstractmethod + def post_deal(self, *args, **kwargs): + """ + Post-processing interface. + """ + pass +``` + +### 接口说明 + +每个Flow对象都需要实现三个核心接口： + +- **prepare**: 用来准备整个workflow的输入数据，设置工作流参数 +- **build_flow**: 用来构建整个workflow的流水线，注册节点和依赖关系 +- **post_deal**: 用来处理workflow的执行结果，转换为对外输出格式 + +### 具体实现示例 + +#### BuildVectorIndexFlow（向量索引构建工作流） + +```python +class BuildVectorIndexFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, texts): + prepared_input.texts = texts + prepared_input.language = "zh" + prepared_input.split_type = "paragraph" + return + + def build_flow(self, texts): + pipeline = GPipeline() + # prepare for workflow input + prepared_input = WkFlowInput() + self.prepare(prepared_input, texts) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + chunk_split_node = ChunkSplitNode() + build_vector_node = BuildVectorIndexNode() + pipeline.registerGElement(chunk_split_node, set(), "chunk_split") + pipeline.registerGElement(build_vector_node, {chunk_split_node}, "build_vector") + + return pipeline + + def post_deal(self, pipeline=None): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + return json.dumps(res, ensure_ascii=False, indent=2) +``` + +#### GraphExtractFlow（图抽取工作流） + +```python +class GraphExtractFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type): + prepared_input.texts = texts + prepared_input.language = "zh" + prepared_input.split_type = "document" + prepared_input.example_prompt = example_prompt + prepared_input.schema = schema + prepare_schema(prepared_input, schema) + return + + def build_flow(self, schema, texts, example_prompt, extract_type): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, schema, texts, example_prompt, extract_type) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + schema_node = SchemaNode() + + chunk_split_node = ChunkSplitNode() + graph_extract_node = ExtractNode() + + pipeline.registerGElement(schema_node, set(), "schema_node") + pipeline.registerGElement(chunk_split_node, set(), "chunk_split") + pipeline.registerGElement( + graph_extract_node, {schema_node, chunk_split_node}, "graph_extract" + ) + + return pipeline + + def post_deal(self, pipeline=None): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + vertices = res.get("vertices", []) + edges = res.get("edges", []) + if not vertices and not edges: + log.info("Please check the schema.(The schema may not match the Doc)") + return json.dumps( + { + "vertices": vertices, + "edges": edges, + "warning": "The schema may not match the Doc", + }, + ensure_ascii=False, + indent=2, + ) + return json.dumps( + {"vertices": vertices, "edges": edges}, + ensure_ascii=False, + indent=2, + ) +``` + +## Node对象设计 + +### 节点生命周期 + +节点以 GNode 为抽象基类，统一生命周期与状态返回。方法职责与约定如下： + +#### 初始化阶段 + +- **init()**: + - **责任**：完成节点级初始化工作（如绑定共享上下文、准备参数区），确保节点具备运行所需的最小环境 + - **约定**：仅做轻量初始化，不执行业务逻辑；返回状态用于判断是否可继续 + +- **node_init()**: + - **责任**：解析与校验本次运行所需的输入（通常来自 wk_input），构建运行期依赖（如内部配置、变换器、资源句柄） + - **约定**：输入缺失或不合法时，应返回错误状态并中止后续执行；不产生对外可见的业务结果 + +#### 运行阶段 + +- **run()**: + - **责任**：执行业务主流程（纯计算或 I/O），在完成后将节点产出写入共享状态（wkflow_state/上下文） + - **约定**： + - 进入前应先调用 node_init() 并检查其返回状态 + - 对共享状态的写操作需遵循并发安全约定（如加锁/解锁） + - 出错使用统一状态返回，不抛出未捕获异常到流程编排层 + +### 输入/输出与上下文约定 + +- **输入**：通过编排层预置于参数区（如 wk_input），节点在 node_init() 中读取并校验 +- **输出**：通过共享状态容器（如 wkflow_state/上下文）对外暴露，键/字段命名应稳定可预期，供下游节点消费 + +### 错误处理约定 + +- 统一以状态对象表示成功/失败与信息；错误应尽早返回，避免在 run() 中继续副作用操作 +- 对可预见的校验类错误使用明确的错误信息，便于定位问题与编排层记录 + +### 并发与可重入约定 + +- 共享状态的写入需在临界区内完成；读取视数据一致性要求决定是否加锁 +- 节点应尽量保持无副作用或将副作用范围收敛在可控区域，以支持重试与复用 + +### 可测试性与解耦 + +- 业务纯逻辑应与框架交互解耦，优先封装为可单测的纯函数/内部方法 +- 节点仅负责生命周期编排与上下文读写，具体策略与算法通过内部可替换组件提供 + +### 节点类型 + +#### 文档处理节点 +- **ChunkSplitNode**: 文档分块处理节点 + - 功能：将输入文档按照指定策略进行分块 + - 输入：原始文档文本 + - 输出：分块后的文档片段 + +#### 索引构建节点 +- **BuildVectorIndexNode**: 向量索引构建节点 + - 功能：基于文档分块构建向量索引 + - 输入：文档分块 + - 输出：向量索引数据 + +#### 模式管理节点 +- **SchemaManagerNode**: 图模式管理节点 + - 功能：从HugeGraph获取图模式信息 + - 输入：图名称 + - 输出：图模式定义 + +- **CheckSchemaNode**: 模式校验节点 + - 功能：校验用户定义的图模式 + - 输入：用户定义的JSON模式 + - 输出：校验后的模式定义 + +#### 图抽取节点 +- **InfoExtractNode**: 信息抽取节点 + - 功能：从文档中抽取三元组信息 + - 输入：文档分块和模式定义 + - 输出：抽取的三元组数据 + +- **PropertyGraphExtractNode**: 属性图抽取节点 + - 功能：从文档中抽取属性图结构 + - 输入：文档分块和模式定义 + - 输出：抽取的顶点和边数据 + +#### 模式构建节点 +- **SchemaBuildNode**: 模式构建节点 + - 功能：基于文档和查询示例构建图模式 + - 输入：文档文本、查询示例、少样本模式 + - 输出：构建的图模式定义 + +#### 提示词生成节点 +- **PromptGenerateNode**: 提示词生成节点 + - 功能：基于源文本、场景和示例名称生成提示词 + - 输入：源文本、场景、示例名称 + - 输出：生成的提示词 + + +## 测试策略 + +### 测试目标 + +目前的测试策略主要目标是保证移植之后的workflow和移植之前的workflow执行结果、程序行为一致。 + +### 测试范围 + +#### 1. 功能测试 +- **工作流执行结果一致性**：确保新架构下的工作流执行结果与原有实现完全一致 +- **输入输出格式验证**：验证输入参数处理和输出格式转换的正确性 +- **错误处理测试**：确保错误场景下的行为与预期一致 + +#### 2. 性能测试 +- **流水线复用效果**：验证流水线复用机制的性能提升效果 +- **并发执行测试**：测试多工作流并发执行的稳定性和性能 +- **资源使用测试**：监控内存和CPU使用情况，确保资源使用合理 + +#### 3. 稳定性测试 +- **长时间运行测试**：验证系统在长时间运行下的稳定性 +- **异常恢复测试**：测试系统在异常情况下的恢复能力 +- **内存泄漏测试**：确保流水线复用不会导致内存泄漏 + +### 测试方法 + +#### 1. 单元测试 +- 对每个Flow类进行单元测试 +- 对每个Node类进行单元测试 +- 对Scheduler调度逻辑进行测试 + +#### 2. 集成测试 +- 端到端工作流测试 +- 多工作流组合测试 +- 与外部系统集成测试 + +#### 3. 性能基准测试 +- 建立性能基准线 +- 对比新旧架构的性能差异 +- 监控关键性能指标 + +### 测试数据 + +#### 1. 标准测试数据集 +- 准备标准化的测试文档 +- 准备标准化的图模式定义 +- 准备标准化的期望输出结果 + +#### 2. 边界测试数据 +- 空输入测试 +- 大文件测试 +- 特殊字符测试 +- 异常格式测试 + +### 测试环境 + +#### 1. 开发环境测试 +- 本地开发环境的功能验证 +- 快速迭代测试 + +#### 2. 测试环境验证 +- 模拟生产环境的完整测试 +- 性能压力测试 + +#### 3. 生产环境验证 +- 灰度发布验证 +- 生产环境监控 + +### 测试自动化 + +#### 1. CI/CD集成 +- 自动化测试流程集成 +- 代码提交触发测试 +- 测试结果自动报告 + +#### 2. 回归测试 +- 定期执行回归测试 +- 确保新功能不影响现有功能 +- 性能回归检测 + +### 测试指标 + +#### 1. 功能指标 +- 测试覆盖率 > 90% +- 功能正确性 100% +- 错误处理覆盖率 > 95% + +#### 2. 性能指标 +- 响应时间提升 > 20% +- 吞吐量提升 > 30% +- 资源使用优化 > 15% + +#### 3. 稳定性指标 +- 系统可用性 > 99.9% +- 平均故障恢复时间 < 5分钟 +- 内存泄漏率 = 0% diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/requirements.md b/.vibedev/spec/hugegraph-llm/fixed_flow/requirements.md new file mode 100644 index 000000000..095027369 --- /dev/null +++ b/.vibedev/spec/hugegraph-llm/fixed_flow/requirements.md @@ -0,0 +1,24 @@ +## 需求列表 + +### 核心框架设计 + +**核心**：Scheduler类中的schedule_flow设计与实现 + +**验收标准**： +1.1. 核心框架尽可能复用资源，避免资源的重复分配和释放 +1.2. 应该保证正常的请求处理指标要求 +1.3. 应该能够配置框架整体使用的资源上限 + +### 固定工作流移植 + +**核心**：移植Web Demo中的所有用例 +2.1. 保证使用核心框架移植后的工作流的程序行为和移植之前保持一致即可 + +**已完成的工作流类型**： +- build_vector_index: 向量索引构建工作流 +- graph_extract: 图抽取工作流 +- import_graph_data: 图数据导入工作流 +- update_vid_embeddings: 向量更新工作流 +- get_graph_index_info: 图索引信息获取工作流 +- build_schema: 模式构建工作流 +- prompt_generate: 提示词生成工作流 diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/tasks.md b/.vibedev/spec/hugegraph-llm/fixed_flow/tasks.md new file mode 100644 index 000000000..a84aee2ff --- /dev/null +++ b/.vibedev/spec/hugegraph-llm/fixed_flow/tasks.md @@ -0,0 +1,36 @@ +# HugeGraph-ai 固定工作流框架设计和用例移植 + +本文档将 HugeGraph 固定工作流框架设计和用例移植转换为一系列可执行的编码任务。 + +## 1. schedule_flow设计与实现 + +- [x] **1.1 构建Scheduler框架1.0** + - 需要能够复用已经创建过的Pipeline(Pipeline Pooling) + - 使用CGraph(Graph-based engine)作为底层执行引擎 + - 不同Node之间松耦合 + +- [ ] **1.2 优化Scheduler框架资源配置** + - 支持用户配置底层线程池参数 + - 现有的workflow可能会根据输入有细小的变化，导致相同的用例得到不同的workflow，怎么解决这个问题呢？ + - Node/Operator解耦，Node负责生命周期和上下文，Operator只关注业务逻辑 + - Flow只负责组装Node，所有业务逻辑下沉到Node/Operator + - Scheduler支持多类型Flow注册，注册方式更灵活 + +- [ ] **1.3 优化Scheduler框架资源使用** + - 根据负载控制每个PipelineManager管理的Pipeline数量，实现动态扩缩容 + - Node层支持参数区自动绑定和并发安全 + - Operator只需实现run(data_json)方法，Node负责调度和结果写回 + +## 2. 固定工作流用例移植 + +- [x] **2.1 build_vector_index workflow移植** +- [x] **2.2 graph_extract workflow移植** +- [x] **2.3 import_graph_data workflow移植** + - 基于Node/Operator机制实现import_graph_data工作流 +- [x] **2.4 update_vid_embeddings workflow移植** + - 基于Node/Operator机制实现update_vid_embeddings工作流 +- [x] **2.5 get_graph_index_info workflow移植** +- [x] **2.6 build_schema workflow移植** + - 基于Node/Operator机制实现build_schema工作流 +- [x] **2.7 prompt_generate workflow移植** + - 基于Node/Operator机制实现prompt_generate工作流 diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 9897f420f..4aa476942 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -26,8 +26,7 @@ from hugegraph_llm.config import huge_settings from hugegraph_llm.config import prompt from hugegraph_llm.config import resource_path -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.llm_op.prompt_generate import PromptGenerate +from hugegraph_llm.flows.scheduler import SchedulerSingleton from hugegraph_llm.utils.graph_index_utils import ( get_graph_index_info, clean_all_graph_index, @@ -61,7 +60,7 @@ def store_prompt(doc, schema, example_prompt): def generate_prompt_for_ui(source_text, scenario, example_name): """ - Handles the UI logic for generating a new prompt. It calls the PromptGenerate operator. + Handles the UI logic for generating a new prompt using the new workflow architecture. """ if not all([source_text, scenario, example_name]): gr.Warning( @@ -69,19 +68,13 @@ def generate_prompt_for_ui(source_text, scenario, example_name): ) return gr.update() try: - prompt_generator = PromptGenerate(llm=LLMs().get_chat_llm()) - context = { - "source_text": source_text, - "scenario": scenario, - "example_name": example_name, - } - result_context = prompt_generator.run(context) - # Presents the result of generating prompt - generated_prompt = result_context.get( - "generated_extract_prompt", "Generation failed. Please check the logs." + # using new architecture + scheduler = SchedulerSingleton.get_instance() + result = scheduler.schedule_flow( + "prompt_generate", source_text, scenario, example_name ) gr.Info("Prompt generated successfully!") - return generated_prompt + return result except Exception as e: log.error("Error generating Prompt: %s", e, exc_info=True) raise gr.Error(f"Error generating Prompt: {e}") from e diff --git a/hugegraph-llm/src/hugegraph_llm/flows/build_schema.py b/hugegraph-llm/src/hugegraph_llm/flows/build_schema.py new file mode 100644 index 000000000..6bbcb8512 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/build_schema.py @@ -0,0 +1,71 @@ +# 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.common import BaseFlow +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.nodes.llm_node.schema_build import SchemaBuildNode +from hugegraph_llm.utils.log import log + +import json +from PyCGraph import GPipeline + + +class BuildSchemaFlow(BaseFlow): + def __init__(self): + pass + + def prepare( + self, + prepared_input: WkFlowInput, + texts=None, + query_examples=None, + few_shot_schema=None, + ): + prepared_input.texts = texts + # Optional fields packed into wk_input for SchemaBuildNode + # Keep raw values; node will parse if strings + prepared_input.query_examples = query_examples + prepared_input.few_shot_schema = few_shot_schema + return + + def build_flow(self, texts=None, query_examples=None, few_shot_schema=None): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare( + prepared_input, + texts=texts, + query_examples=query_examples, + few_shot_schema=few_shot_schema, + ) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + schema_build_node = SchemaBuildNode() + pipeline.registerGElement(schema_build_node, set(), "schema_build") + + return pipeline + + def post_deal(self, pipeline=None): + state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + if "schema" not in state_json: + return "" + res = state_json["schema"] + try: + formatted_schema = json.dumps(res, ensure_ascii=False, indent=2) + return formatted_schema + except (TypeError, ValueError) as e: + log.error("Failed to format schema: %s", e) + return str(res) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py index f1ee8c1c4..9a07b5dba 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py @@ -14,13 +14,13 @@ # limitations under the License. from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.document_node.chunk_split import ChunkSplitNode +from hugegraph_llm.nodes.index_node.build_vector_index import BuildVectorIndexNode from hugegraph_llm.state.ai_state import WkFlowInput import json from PyCGraph import GPipeline -from hugegraph_llm.operators.document_op.chunk_split import ChunkSplitNode -from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndexNode from hugegraph_llm.state.ai_state import WkFlowState diff --git a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py new file mode 100644 index 000000000..fa10d0199 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py @@ -0,0 +1,68 @@ +# 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 os + +from hugegraph_llm.config import huge_settings, llm_settings, resource_path +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.indices.vector_index import VectorIndex +from hugegraph_llm.models.embeddings.init_embedding import model_map +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.nodes.hugegraph_node.fetch_graph_data import FetchGraphDataNode +from PyCGraph import GPipeline +from hugegraph_llm.utils.embedding_utils import ( + get_filename_prefix, + get_index_folder_name, +) + + +class GetGraphIndexInfoFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, *args, **kwargs): + return + + def build_flow(self, *args, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, *args, **kwargs) + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + fetch_node = FetchGraphDataNode() + pipeline.registerGElement(fetch_node, set(), "fetch_node") + return pipeline + + def post_deal(self, pipeline=None): + graph_summary_info = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, + model_map.get(llm_settings.embedding_type, None), + ) + try: + vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) + except FileNotFoundError: + return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) + graph_summary_info["vid_index"] = { + "embed_dim": vector_index.index.d, + "num_vectors": vector_index.index.ntotal, + "num_vids": len(vector_index.properties), + } + return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py index f1a6c5f6f..1b0c98253 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py @@ -16,14 +16,10 @@ import json from PyCGraph import GPipeline from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.document_node.chunk_split import ChunkSplitNode +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.nodes.llm_node.extract_info import ExtractNode from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState -from hugegraph_llm.operators.common_op.check_schema import CheckSchemaNode -from hugegraph_llm.operators.document_op.chunk_split import ChunkSplitNode -from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManagerNode -from hugegraph_llm.operators.llm_op.info_extract import InfoExtractNode -from hugegraph_llm.operators.llm_op.property_graph_extract import ( - PropertyGraphExtractNode, -) from hugegraph_llm.utils.log import log @@ -31,21 +27,6 @@ class GraphExtractFlow(BaseFlow): def __init__(self): pass - def _import_schema( - self, - from_hugegraph=None, - from_extraction=None, - from_user_defined=None, - ): - if from_hugegraph: - return SchemaManagerNode() - elif from_user_defined: - return CheckSchemaNode() - elif from_extraction: - raise NotImplementedError("Not implemented yet") - else: - raise ValueError("No input data / invalid schema type") - def prepare( self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type ): @@ -55,17 +36,7 @@ def prepare( prepared_input.split_type = "document" prepared_input.example_prompt = example_prompt prepared_input.schema = schema - schema = schema.strip() - if schema.startswith("{"): - try: - schema = json.loads(schema) - prepared_input.schema = schema - except json.JSONDecodeError as exc: - log.error("Invalid JSON format in schema. Please check it again.") - raise ValueError("Invalid JSON format in schema.") from exc - else: - log.info("Get schema '%s' from graphdb.", schema) - prepared_input.graph_name = schema + prepared_input.extract_type = extract_type return def build_flow(self, schema, texts, example_prompt, extract_type): @@ -76,27 +47,10 @@ def build_flow(self, schema, texts, example_prompt, extract_type): pipeline.createGParam(prepared_input, "wkflow_input") pipeline.createGParam(WkFlowState(), "wkflow_state") - schema = schema.strip() - schema_node = None - if schema.startswith("{"): - try: - schema = json.loads(schema) - schema_node = self._import_schema(from_user_defined=schema) - except json.JSONDecodeError as exc: - log.error("Invalid JSON format in schema. Please check it again.") - raise ValueError("Invalid JSON format in schema.") from exc - else: - log.info("Get schema '%s' from graphdb.", schema) - schema_node = self._import_schema(from_hugegraph=schema) + schema_node = SchemaNode() chunk_split_node = ChunkSplitNode() - graph_extract_node = None - if extract_type == "triples": - graph_extract_node = InfoExtractNode() - elif extract_type == "property_graph": - graph_extract_node = PropertyGraphExtractNode() - else: - raise ValueError(f"Unsupported extract_type: {extract_type}") + graph_extract_node = ExtractNode() pipeline.registerGElement(schema_node, set(), "schema_node") pipeline.registerGElement(chunk_split_node, set(), "chunk_split") pipeline.registerGElement( diff --git a/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py b/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py new file mode 100644 index 000000000..5581ef107 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py @@ -0,0 +1,65 @@ +# 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 gradio as gr +from PyCGraph import GPipeline +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.hugegraph_node.commit_to_hugegraph import Commit2GraphNode +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.utils.log import log + + +class ImportGraphDataFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, data, schema): + try: + data_json = json.loads(data.strip()) if isinstance(data, str) else data + except json.JSONDecodeError as e: + raise ValueError(f"Invalid JSON for 'data': {e.msg}") from e + log.debug( + "Import graph data (truncated): %s", + (data[:512] + "...") + if isinstance(data, str) and len(data) > 512 + else (data if isinstance(data, str) else ""), + ) + prepared_input.data_json = data_json + prepared_input.schema = schema + return + + def build_flow(self, data, schema): + pipeline = GPipeline() + prepared_input = WkFlowInput() + # prepare input data + self.prepare(prepared_input, data, schema) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + schema_node = SchemaNode() + commit_node = Commit2GraphNode() + pipeline.registerGElement(schema_node, set(), "schema_node") + pipeline.registerGElement(commit_node, {schema_node}, "commit_node") + + return pipeline + + def post_deal(self, pipeline=None): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + gr.Info("Import graph data successfully!") + return json.dumps(res, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py new file mode 100644 index 000000000..aece6bd61 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py @@ -0,0 +1,63 @@ +# 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.common import BaseFlow +from hugegraph_llm.nodes.llm_node.prompt_generate import PromptGenerateNode +from hugegraph_llm.state.ai_state import WkFlowInput + +from PyCGraph import GPipeline + +from hugegraph_llm.state.ai_state import WkFlowState + + +class PromptGenerateFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, source_text, scenario, example_name): + """ + Prepare input data for PromptGenerate workflow + """ + prepared_input.source_text = source_text + prepared_input.scenario = scenario + prepared_input.example_name = example_name + return + + def build_flow(self, source_text, scenario, example_name): + """ + Build the PromptGenerate workflow + """ + pipeline = GPipeline() + # Prepare workflow input + prepared_input = WkFlowInput() + self.prepare(prepared_input, source_text, scenario, example_name) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + # Create PromptGenerate node + prompt_generate_node = PromptGenerateNode() + pipeline.registerGElement(prompt_generate_node, set(), "prompt_generate") + + return pipeline + + def post_deal(self, pipeline=None): + """ + Process the execution result of PromptGenerate workflow + """ + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + return res.get( + "generated_extract_prompt", "Generation failed. Please check the logs." + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index b096310db..559540ce3 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -15,10 +15,15 @@ import threading from typing import Dict, Any -from PyCGraph import GPipelineManager +from PyCGraph import GPipeline, GPipelineManager from hugegraph_llm.flows.build_vector_index import BuildVectorIndexFlow from hugegraph_llm.flows.common import BaseFlow from hugegraph_llm.flows.graph_extract import GraphExtractFlow +from hugegraph_llm.flows.import_graph_data import ImportGraphDataFlow +from hugegraph_llm.flows.update_vid_embeddings import UpdateVidEmbeddingsFlows +from hugegraph_llm.flows.get_graph_index_info import GetGraphIndexInfoFlow +from hugegraph_llm.flows.build_schema import BuildSchemaFlow +from hugegraph_llm.flows.prompt_generate import PromptGenerateFlow from hugegraph_llm.utils.log import log @@ -37,6 +42,26 @@ def __init__(self, max_pipeline: int = 10): "manager": GPipelineManager(), "flow": GraphExtractFlow(), } + self.pipeline_pool["import_graph_data"] = { + "manager": GPipelineManager(), + "flow": ImportGraphDataFlow(), + } + self.pipeline_pool["update_vid_embeddings"] = { + "manager": GPipelineManager(), + "flow": UpdateVidEmbeddingsFlows(), + } + self.pipeline_pool["get_graph_index_info"] = { + "manager": GPipelineManager(), + "flow": GetGraphIndexInfoFlow(), + } + self.pipeline_pool["build_schema"] = { + "manager": GPipelineManager(), + "flow": BuildSchemaFlow(), + } + self.pipeline_pool["prompt_generate"] = { + "manager": GPipelineManager(), + "flow": PromptGenerateFlow(), + } self.max_pipeline = max_pipeline # TODO: Implement Agentic Workflow @@ -46,9 +71,9 @@ def agentic_flow(self): def schedule_flow(self, flow: str, *args, **kwargs): if flow not in self.pipeline_pool: raise ValueError(f"Unsupported workflow {flow}") - manager = self.pipeline_pool[flow]["manager"] + manager: GPipelineManager = self.pipeline_pool[flow]["manager"] flow: BaseFlow = self.pipeline_pool[flow]["flow"] - pipeline = manager.fetch() + pipeline: GPipeline = manager.fetch() if pipeline is None: # call coresponding flow_func to create new workflow pipeline = flow.build_flow(*args, **kwargs) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py b/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py new file mode 100644 index 000000000..b3f0d9923 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py @@ -0,0 +1,47 @@ +# 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 PyCGraph import CStatus, GPipeline +from hugegraph_llm.flows.common import BaseFlow, WkFlowInput +from hugegraph_llm.nodes.hugegraph_node.fetch_graph_data import FetchGraphDataNode +from hugegraph_llm.nodes.index_node.build_semantic_index import BuildSemanticIndexNode +from hugegraph_llm.state.ai_state import WkFlowState + + +class UpdateVidEmbeddingsFlows(BaseFlow): + def prepare(self, prepared_input: WkFlowInput): + return CStatus() + + def build_flow(self): + pipeline = GPipeline() + prepared_input = WkFlowInput() + # prepare input data + self.prepare(prepared_input) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + fetch_node = FetchGraphDataNode() + build_node = BuildSemanticIndexNode() + pipeline.registerGElement(fetch_node, set(), "fetch_node") + pipeline.registerGElement(build_node, {fetch_node}, "build_node") + + return pipeline + + def post_deal(self, pipeline): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + removed_num = res.get("removed_vid_vector_num", 0) + added_num = res.get("added_vid_vector_num", 0) + return f"Removed {removed_num} vectors, added {added_num} vectors." diff --git a/hugegraph-llm/src/hugegraph_llm/flows/utils.py b/hugegraph-llm/src/hugegraph_llm/flows/utils.py new file mode 100644 index 000000000..b4ba05c84 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/utils.py @@ -0,0 +1,34 @@ +# 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 hugegraph_llm.state.ai_state import WkFlowInput +from hugegraph_llm.utils.log import log + + +def prepare_schema(prepared_input: WkFlowInput, schema): + schema = schema.strip() + if schema.startswith("{"): + try: + schema = json.loads(schema) + prepared_input.schema = schema + except json.JSONDecodeError as exc: + log.error("Invalid JSON format in schema. Please check it again.") + raise ValueError("Invalid JSON format in schema.") from exc + else: + log.info("Get schema '%s' from graphdb.", schema) + prepared_input.graph_name = schema + return diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py new file mode 100644 index 000000000..0ea0675c0 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py @@ -0,0 +1,71 @@ +# 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 PyCGraph import GNode, CStatus +from hugegraph_llm.nodes.util import init_context +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class BaseNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + return init_context(self) + + def node_init(self): + """ + Node initialization method, can be overridden by subclasses. + Returns a CStatus object indicating whether initialization succeeded. + """ + return CStatus() + + def run(self): + """ + Main logic for node execution, can be overridden by subclasses. + Returns a CStatus object indicating whether execution succeeded. + """ + sts = self.node_init() + if sts.isErr(): + return sts + self.context.lock() + try: + data_json = self.context.to_json() + finally: + self.context.unlock() + + try: + res = self.operator_schedule(data_json) + except Exception as exc: + import traceback + + node_info = f"Node type: {type(self).__name__}, Node object: {self}" + err_msg = f"Node failed: {exc}\n{node_info}\n{traceback.format_exc()}" + return CStatus(-1, err_msg) + + self.context.lock() + try: + if isinstance(res, dict): + self.context.assign_from_json(res) + finally: + self.context.unlock() + return CStatus() + + def operator_schedule(self, data_json): + """ + Interface for scheduling the operator, can be overridden by subclasses. + Returns a CStatus object indicating whether scheduling succeeded. + """ + pass 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 new file mode 100644 index 000000000..4c5acbe97 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/document_node/chunk_split.py @@ -0,0 +1,43 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.operators.document_op.chunk_split import ChunkSplit +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class ChunkSplitNode(BaseNode): + chunk_split_op: ChunkSplit + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + if ( + self.wk_input.texts is None + or self.wk_input.language is None + or self.wk_input.split_type is None + ): + return CStatus(-1, "Error occurs when prepare for workflow input") + texts = self.wk_input.texts + language = self.wk_input.language + split_type = self.wk_input.split_type + if isinstance(texts, str): + texts = [texts] + self.chunk_split_op = ChunkSplit(texts, split_type, language) + return CStatus() + + def operator_schedule(self, data_json): + 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 new file mode 100644 index 000000000..b576e8170 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/commit_to_hugegraph.py @@ -0,0 +1,35 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph import Commit2Graph +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class Commit2GraphNode(BaseNode): + commit_to_graph_op: Commit2Graph + context: WkFlowState = None + wk_input: WkFlowInput = None + + 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() + return CStatus() + + def operator_schedule(self, data_json): + return self.commit_to_graph_op.run(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 new file mode 100644 index 000000000..b2434e524 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/fetch_graph_data.py @@ -0,0 +1,33 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.hugegraph_op.fetch_graph_data import FetchGraphData +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.utils.hugegraph_utils import get_hg_client + + +class FetchGraphDataNode(BaseNode): + fetch_graph_data_op: FetchGraphData + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + self.fetch_graph_data_op = FetchGraphData(get_hg_client()) + return CStatus() + + def operator_schedule(self, data_json): + return self.fetch_graph_data_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py new file mode 100644 index 000000000..71c490b20 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.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 json + +from PyCGraph import CStatus +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.common_op.check_schema import CheckSchema +from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.utils.log import log + + +class SchemaNode(BaseNode): + schema_manager: SchemaManager + check_schema: CheckSchema + context: WkFlowState = None + wk_input: WkFlowInput = None + + schema = None + + def _import_schema( + self, + from_hugegraph=None, + from_extraction=None, + from_user_defined=None, + ): + if from_hugegraph: + return SchemaManager(from_hugegraph) + elif from_user_defined: + return CheckSchema(from_user_defined) + elif from_extraction: + raise NotImplementedError("Not implemented yet") + else: + raise ValueError("No input data / invalid schema type") + + def node_init(self): + self.schema = self.wk_input.schema + self.schema = self.schema.strip() + if self.schema.startswith("{"): + try: + schema = json.loads(self.schema) + self.check_schema = self._import_schema(from_user_defined=schema) + except json.JSONDecodeError as exc: + log.error("Invalid JSON format in schema. Please check it again.") + raise ValueError("Invalid JSON format in schema.") from exc + else: + log.info("Get schema '%s' from graphdb.", self.schema) + self.schema_manager = self._import_schema(from_hugegraph=self.schema) + return CStatus() + + def operator_schedule(self, data_json): + print(f"check data json {data_json}") + if self.schema.startswith("{"): + try: + return self.check_schema.run(data_json) + except json.JSONDecodeError as exc: + log.error("Invalid JSON format in schema. Please check it again.") + raise ValueError("Invalid JSON format in schema.") from exc + else: + log.info("Get schema '%s' from graphdb.", self.schema) + return self.schema_manager.run(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 new file mode 100644 index 000000000..ab31fa394 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_semantic_index.py @@ -0,0 +1,34 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.config import llm_settings +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class BuildSemanticIndexNode(BaseNode): + build_semantic_index_op: BuildSemanticIndex + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + self.build_semantic_index_op = BuildSemanticIndex(get_embedding(llm_settings)) + return CStatus() + + def operator_schedule(self, data_json): + return self.build_semantic_index_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py new file mode 100644 index 000000000..cf2f9b677 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py @@ -0,0 +1,34 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.config import llm_settings +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class BuildVectorIndexNode(BaseNode): + build_vector_index_op: BuildVectorIndex + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + self.build_vector_index_op = BuildVectorIndex(get_embedding(llm_settings)) + return CStatus() + + def operator_schedule(self, data_json): + return self.build_vector_index_op.run(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 new file mode 100644 index 000000000..8bceed804 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/extract_info.py @@ -0,0 +1,52 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.config import llm_settings +from hugegraph_llm.models.llms.init_llm import get_chat_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 +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class ExtractNode(BaseNode): + property_graph_extract: PropertyGraphExtract + info_extract: InfoExtract + context: WkFlowState = None + wk_input: WkFlowInput = None + + extract_type: str = None + + def node_init(self): + llm = get_chat_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 + extract_type = self.wk_input.extract_type + self.extract_type = extract_type + if extract_type == "triples": + self.info_extract = InfoExtract(llm, example_prompt) + elif extract_type == "property_graph": + self.property_graph_extract = PropertyGraphExtract(llm, example_prompt) + else: + return CStatus(-1, f"Unsupported extract_type: {extract_type}") + return CStatus() + + def operator_schedule(self, data_json): + if self.extract_type == "triples": + return self.info_extract.run(data_json) + elif self.extract_type == "property_graph": + return self.property_graph_extract.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py new file mode 100644 index 000000000..317f9e6ac --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py @@ -0,0 +1,59 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.config import llm_settings +from hugegraph_llm.models.llms.init_llm import get_chat_llm +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.llm_op.prompt_generate import PromptGenerate +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class PromptGenerateNode(BaseNode): + prompt_generate: PromptGenerate + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + """ + Node initialization method, initialize PromptGenerate operator + """ + llm = get_chat_llm(llm_settings) + if not all( + [ + self.wk_input.source_text, + self.wk_input.scenario, + self.wk_input.example_name, + ] + ): + return CStatus( + -1, + "Missing required parameters: source_text, scenario, or example_name", + ) + + self.prompt_generate = PromptGenerate(llm) + context = { + "source_text": self.wk_input.source_text, + "scenario": self.wk_input.scenario, + "example_name": self.wk_input.example_name, + } + self.context.assign_from_json(context) + return CStatus() + + def operator_schedule(self, data_json): + """ + Schedule the execution of PromptGenerate operator + """ + return self.prompt_generate.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py new file mode 100644 index 000000000..a28b41346 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py @@ -0,0 +1,91 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.models.llms.init_llm import get_chat_llm +from hugegraph_llm.config import llm_settings +from hugegraph_llm.operators.llm_op.schema_build import SchemaBuilder +from hugegraph_llm.utils.log import log + + +class SchemaBuildNode(BaseNode): + schema_builder: SchemaBuilder + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + llm = get_chat_llm(llm_settings) + self.schema_builder = SchemaBuilder(llm) + + # texts -> raw_texts + raw_texts = [] + if self.wk_input.texts: + if isinstance(self.wk_input.texts, list): + raw_texts = [t for t in self.wk_input.texts if isinstance(t, str)] + elif isinstance(self.wk_input.texts, str): + raw_texts = [self.wk_input.texts] + + # query_examples: already parsed list[dict] or raw JSON string + query_examples = [] + qe_src = self.wk_input.query_examples if self.wk_input.query_examples else None + if qe_src: + try: + parsed_examples = json.loads(qe_src) + # Validate and retain the description and gremlin fields + query_examples = [ + { + "description": ex.get("description", ""), + "gremlin": ex.get("gremlin", ""), + } + for ex in parsed_examples + if isinstance(ex, dict) and "description" in ex and "gremlin" in ex + ] + except json.JSONDecodeError as e: + return CStatus(-1, f"Query Examples is not in a valid JSON format: {e}") + + # few_shot_schema: already parsed dict or raw JSON string + few_shot_schema = {} + fss_src = ( + self.wk_input.few_shot_schema if self.wk_input.few_shot_schema else None + ) + if fss_src: + try: + few_shot_schema = json.loads(fss_src) + except json.JSONDecodeError as e: + return CStatus( + -1, f"Few Shot Schema is not in a valid JSON format: {e}" + ) + + _context_payload = { + "raw_texts": raw_texts, + "query_examples": query_examples, + "few_shot_schema": few_shot_schema, + } + self.context.assign_from_json(_context_payload) + + return CStatus() + + def operator_schedule(self, data_json): + try: + schema_result = self.schema_builder.run(data_json) + + return {"schema": schema_result} + except Exception as e: + log.error("Failed to generate schema: %s", e) + return {"schema": f"Schema generation failed: {e}"} diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/util.py b/hugegraph-llm/src/hugegraph_llm/nodes/util.py new file mode 100644 index 000000000..60bdc2e86 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/util.py @@ -0,0 +1,27 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from PyCGraph import CStatus + + +def init_context(obj) -> CStatus: + try: + obj.context = obj.getGParamWithNoEmpty("wkflow_state") + obj.wk_input = obj.getGParamWithNoEmpty("wkflow_input") + if obj.context is None or obj.wk_input is None: + return CStatus(-1, "Required workflow parameters not found") + return CStatus() + except Exception as e: + return CStatus(-1, f"Failed to initialize context: {str(e)}") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index 7a533517a..c1c742032 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -20,12 +20,8 @@ from hugegraph_llm.enums.property_cardinality import PropertyCardinality from hugegraph_llm.enums.property_data_type import PropertyDataType -from hugegraph_llm.operators.util import init_context from hugegraph_llm.utils.log import log -from PyCGraph import GNode, CStatus -from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState - def log_and_raise(message: str) -> None: log.warning(message) @@ -174,159 +170,3 @@ def _add_missing_properties( } ) property_label_set.add(prop) - - -class CheckSchemaNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None - - def init(self): - return init_context(self) - - def node_init(self): - if self.wk_input.schema is None: - return CStatus(-1, "Error occurs when prepare for workflow input") - self.data = self.wk_input.schema - return CStatus() - - def run(self) -> CStatus: - # init workflow input - sts = self.node_init() - if sts.isErr(): - return sts - # 1. Validate the schema structure - self.context.lock() - schema = self.data or self.context.schema - self._validate_schema(schema) - # 2. Process property labels and also create a set for it - property_labels, property_label_set = self._process_property_labels(schema) - # 3. Process properties in given vertex/edge labels - self._process_vertex_labels(schema, property_labels, property_label_set) - self._process_edge_labels(schema, property_labels, property_label_set) - # 4. Update schema with processed pks - schema["propertykeys"] = property_labels - self.context.schema = schema - self.context.unlock() - return CStatus() - - def _validate_schema(self, schema: Dict[str, Any]) -> None: - check_type(schema, dict, "Input data is not a dictionary.") - if "vertexlabels" not in schema or "edgelabels" not in schema: - log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") - check_type( - schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." - ) - check_type( - schema["edgelabels"], list, "'edgelabels' in input data is not a list." - ) - - def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): - property_labels = schema.get("propertykeys", []) - check_type( - property_labels, - list, - "'propertykeys' in input data is not of correct type.", - ) - property_label_set = {label["name"] for label in property_labels} - return property_labels, property_label_set - - def _process_vertex_labels( - self, schema: Dict[str, Any], property_labels: list, property_label_set: set - ) -> None: - for vertex_label in schema["vertexlabels"]: - self._validate_vertex_label(vertex_label) - properties = vertex_label["properties"] - primary_keys = self._process_keys( - vertex_label, "primary_keys", properties[:1] - ) - if len(primary_keys) == 0: - log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") - vertex_label["primary_keys"] = primary_keys - nullable_keys = self._process_keys( - vertex_label, "nullable_keys", properties[1:] - ) - vertex_label["nullable_keys"] = nullable_keys - self._add_missing_properties( - properties, property_labels, property_label_set - ) - - def _process_edge_labels( - self, schema: Dict[str, Any], property_labels: list, property_label_set: set - ) -> None: - for edge_label in schema["edgelabels"]: - self._validate_edge_label(edge_label) - properties = edge_label.get("properties", []) - self._add_missing_properties( - properties, property_labels, property_label_set - ) - - def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: - check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") - if "name" not in vertex_label: - log_and_raise("VertexLabel in input data does not contain 'name'.") - check_type( - vertex_label["name"], str, "'name' in vertex_label is not of correct type." - ) - if "properties" not in vertex_label: - log_and_raise("VertexLabel in input data does not contain 'properties'.") - check_type( - vertex_label["properties"], - list, - "'properties' in vertex_label is not of correct type.", - ) - if len(vertex_label["properties"]) == 0: - log_and_raise("'properties' in vertex_label is empty.") - - def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: - check_type(edge_label, dict, "EdgeLabel in input data is not a dictionary.") - if ( - "name" not in edge_label - or "source_label" not in edge_label - or "target_label" not in edge_label - ): - log_and_raise( - "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." - ) - check_type( - edge_label["name"], str, "'name' in edge_label is not of correct type." - ) - check_type( - edge_label["source_label"], - str, - "'source_label' in edge_label is not of correct type.", - ) - check_type( - edge_label["target_label"], - str, - "'target_label' in edge_label is not of correct type.", - ) - - def _process_keys( - self, label: Dict[str, Any], key_type: str, default_keys: list - ) -> list: - keys = label.get(key_type, default_keys) - check_type( - keys, list, f"'{key_type}' in {label['name']} is not of correct type." - ) - new_keys = [key for key in keys if key in label["properties"]] - return new_keys - - def _add_missing_properties( - self, properties: list, property_labels: list, property_label_set: set - ) -> None: - for prop in properties: - if prop not in property_label_set: - property_labels.append( - { - "name": prop, - "data_type": PropertyDataType.DEFAULT.value, - "cardinality": PropertyCardinality.DEFAULT.value, - } - ) - property_label_set.add(prop) - - def get_result(self): - self.context.lock() - res = self.context.to_json() - self.context.unlock() - return res 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 d779a40ab..c31e77af7 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 @@ -19,8 +19,6 @@ from typing import Literal, Dict, Any, Optional, Union, List from langchain_text_splitters import RecursiveCharacterTextSplitter -from hugegraph_llm.operators.util import init_context -from PyCGraph import GNode, CStatus # Constants LANGUAGE_ZH = "zh" @@ -30,62 +28,6 @@ SPLIT_TYPE_SENTENCE = "sentence" -class ChunkSplitNode(GNode): - def init(self): - return init_context(self) - - def node_init(self): - if ( - self.wk_input.texts is None - or self.wk_input.language is None - or self.wk_input.split_type is None - ): - return CStatus(-1, "Error occurs when prepare for workflow input") - texts = self.wk_input.texts - language = self.wk_input.language - split_type = self.wk_input.split_type - if isinstance(texts, str): - texts = [texts] - self.texts = texts - self.separators = self._get_separators(language) - self.text_splitter = self._get_text_splitter(split_type) - return CStatus() - - def _get_separators(self, language: str) -> List[str]: - if language == LANGUAGE_ZH: - return ["\n\n", "\n", "。", "，", ""] - if language == LANGUAGE_EN: - return ["\n\n", "\n", ".", ",", " ", ""] - raise ValueError("language must be zh or en") - - 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 - if split_type == SPLIT_TYPE_SENTENCE: - return RecursiveCharacterTextSplitter( - chunk_size=50, chunk_overlap=0, separators=self.separators - ).split_text - raise ValueError("Type must be document, paragraph or sentence") - - def run(self): - sts = self.node_init() - if sts.isErr(): - return sts - all_chunks = [] - for text in self.texts: - chunks = self.text_splitter(text) - all_chunks.extend(chunks) - - self.context.lock() - self.context.chunks = all_chunks - self.context.unlock() - return CStatus() - - class ChunkSplit: def __init__( self, 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 5cc846d21..9eec04f7f 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 @@ -40,15 +40,19 @@ def run(self, data: dict) -> Dict[str, Any]: schema = data.get("schema") vertices = data.get("vertices", []) edges = data.get("edges", []) - + print(f"get schema {schema}") if not vertices and not edges: - log.critical("(Loading) Both vertices and edges are empty. Please check the input data again.") + log.critical( + "(Loading) Both vertices and edges are empty. Please check the input data again." + ) raise ValueError("Both vertices and edges input are empty.") if not schema: # TODO: ensure the function works correctly (update the logic later) self.schema_free_mode(data.get("triples", [])) - log.warning("Using schema_free mode, could try schema_define mode for better effect!") + log.warning( + "Using schema_free mode, could try schema_define mode for better effect!" + ) else: self.init_schema_if_need(schema) self.load_into_graph(vertices, edges, schema) @@ -64,7 +68,9 @@ def _set_default_property(self, key, input_properties, property_label_map): # list or set default_value = [] input_properties[key] = default_value - log.warning("Property '%s' missing in vertex, set to '%s' for now", key, default_value) + log.warning( + "Property '%s' missing in vertex, set to '%s' for now", key, default_value + ) def _handle_graph_creation(self, func, *args, **kwargs): try: @@ -78,29 +84,42 @@ 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) - vertex_label_map = {v_label["name"]: v_label for v_label in schema["vertexlabels"]} + 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"]} + property_label_map = { + p_label["name"]: p_label for p_label in schema["propertykeys"] + } for vertex in vertices: input_label = vertex["label"] # 1. ensure the input_label in the graph schema if input_label not in vertex_label_map: - log.critical("(Input) VertexLabel %s not found in schema, skip & need check it!", input_label) + log.critical( + "(Input) VertexLabel %s not found in schema, skip & need check it!", + input_label, + ) continue input_properties = vertex["properties"] vertex_label = vertex_label_map[input_label] primary_keys = vertex_label["primary_keys"] nullable_keys = vertex_label.get("nullable_keys", []) - non_null_keys = [key for key in vertex_label["properties"] if key not in nullable_keys] + non_null_keys = [ + key for key in vertex_label["properties"] if key not in nullable_keys + ] has_problem = False # 2. Handle primary-keys mode vertex for pk in primary_keys: if not input_properties.get(pk): if len(primary_keys) == 1: - log.error("Primary-key '%s' missing in vertex %s, skip it & need check it again", pk, vertex) + log.error( + "Primary-key '%s' missing in vertex %s, skip it & need check it again", + pk, + vertex, + ) has_problem = True break # TODO: transform to Enum first (better in earlier step) @@ -110,14 +129,20 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- input_properties[pk] = default_value_map(data_type) else: input_properties[pk] = [] - log.warning("Primary-key '%s' missing in vertex %s, mark empty & need check it again!", pk, vertex) + log.warning( + "Primary-key '%s' missing in vertex %s, mark empty & need check it again!", + pk, + vertex, + ) if has_problem: continue # 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) + 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(): @@ -125,14 +150,19 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- 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) + log.error( + "Property type/format '%s' is not correct, skip it & need check it again", + key, + ) has_problem = True break if has_problem: continue # TODO: we could try batch add vertices first, setback to single-mode if failed - vid = self._handle_graph_creation(self.client.graph().addVertex, input_label, input_properties).id + vid = self._handle_graph_creation( + self.client.graph().addVertex, input_label, input_properties + ).id vertex["id"] = vid for edge in edges: @@ -142,11 +172,16 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- properties = edge["properties"] if label not in edge_label_map: - log.critical("(Input) EdgeLabel %s not found in schema, skip & need check it!", label) + log.critical( + "(Input) EdgeLabel %s not found in schema, skip & need check it!", + label, + ) 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) + self._handle_graph_creation( + self.client.graph().addEdge, label, start, end, properties + ) def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -170,19 +205,27 @@ def init_schema_if_need(self, schema: dict): source_vertex_label = edge["source_label"] target_vertex_label = edge["target_label"] properties = edge["properties"] - self.schema.edgeLabel(edge_label).sourceLabel(source_vertex_label).targetLabel( - target_vertex_label - ).properties(*properties).nullableKeys(*properties).ifNotExist().create() + self.schema.edgeLabel(edge_label).sourceLabel( + source_vertex_label + ).targetLabel(target_vertex_label).properties(*properties).nullableKeys( + *properties + ).ifNotExist().create() def schema_free_mode(self, 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( + self.schema.vertexLabel("vertex").useCustomizeStringId().properties( "name" ).ifNotExist().create() + self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel( + "vertex" + ).properties("name").ifNotExist().create() - self.schema.indexLabel("vertexByName").onV("vertex").by("name").secondary().ifNotExist().create() - self.schema.indexLabel("edgeByName").onE("edge").by("name").secondary().ifNotExist().create() + 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: s, p, o = (element.strip() for element in item) @@ -196,8 +239,12 @@ def _create_property(self, prop: dict): data_type = PropertyDataType(prop["data_type"]) cardinality = PropertyCardinality(prop["cardinality"]) except ValueError: - log.critical("Invalid data type %s / cardinality %s for property %s, skip & should check it again", - prop["data_type"], prop["cardinality"], name) + log.critical( + "Invalid data type %s / cardinality %s for property %s, skip & should check it again", + prop["data_type"], + prop["cardinality"], + name, + ) return property_key = self.schema.propertyKey(name) @@ -231,7 +278,9 @@ def _set_property_data_type(self, property_key, data_type): log.warning("UUID type is not supported, use text instead") property_key.asText() else: - log.error("Unknown data type %s for property_key %s", data_type, property_key) + log.error( + "Unknown data type %s for property_key %s", data_type, property_key + ) def _set_property_cardinality(self, property_key, cardinality): if cardinality == PropertyCardinality.SINGLE: @@ -241,10 +290,17 @@ def _set_property_cardinality(self, property_key, cardinality): elif cardinality == PropertyCardinality.SET: property_key.valueSet() else: - 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): + 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) @@ -259,14 +315,21 @@ def _check_collection_data_type(self, data_type: str, value) -> bool: def _check_single_data_type(self, data_type: str, value) -> bool: if data_type == PropertyDataType.BOOLEAN.value: return isinstance(value, bool) - if data_type in (PropertyDataType.BYTE.value, PropertyDataType.INT.value, PropertyDataType.LONG.value): + if data_type in ( + PropertyDataType.BYTE.value, + PropertyDataType.INT.value, + PropertyDataType.LONG.value, + ): return isinstance(value, int) if data_type in (PropertyDataType.FLOAT.value, PropertyDataType.DOUBLE.value): return isinstance(value, float) if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below - if data_type == PropertyDataType.DATE.value: # the format should be "yyyy-MM-dd" + if ( + data_type == PropertyDataType.DATE.value + ): # the format should be "yyyy-MM-dd" import re - return isinstance(value, str) and re.match(r'^\d{4}-\d{2}-\d{2}$', value) + + return isinstance(value, str) and re.match(r"^\d{4}-\d{2}-\d{2}$", value) raise ValueError(f"Unknown/Unsupported data type: {data_type}") 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 670c18b4a..c4e2124c3 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 @@ -17,12 +17,8 @@ from typing import Dict, Any, Optional from hugegraph_llm.config import huge_settings -from hugegraph_llm.operators.util import init_context -from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState from pyhugegraph.client import PyHugeClient -from PyCGraph import GNode, CStatus - class SchemaManager: def __init__(self, graph_name: str): @@ -74,74 +70,3 @@ def run(self, context: Optional[Dict[str, Any]]) -> Dict[str, Any]: # TODO: enhance the logic here context["simple_schema"] = self.simple_schema(schema) return context - - -class SchemaManagerNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None - - def init(self): - return init_context(self) - - def node_init(self): - if self.wk_input.graph_name is None: - return CStatus(-1, "Error occurs when prepare for workflow input") - graph_name = self.wk_input.graph_name - self.graph_name = graph_name - self.client = PyHugeClient( - url=huge_settings.graph_url, - graph=self.graph_name, - user=huge_settings.graph_user, - pwd=huge_settings.graph_pwd, - graphspace=huge_settings.graph_space, - ) - self.schema = self.client.schema() - return CStatus() - - def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: - mini_schema = {} - - # Add necessary vertexlabels items (3) - if "vertexlabels" in schema: - mini_schema["vertexlabels"] = [] - for vertex in schema["vertexlabels"]: - new_vertex = { - key: vertex[key] - for key in ["id", "name", "properties"] - if key in vertex - } - mini_schema["vertexlabels"].append(new_vertex) - - # Add necessary edgelabels items (4) - if "edgelabels" in schema: - mini_schema["edgelabels"] = [] - for edge in schema["edgelabels"]: - new_edge = { - key: edge[key] - for key in ["name", "source_label", "target_label", "properties"] - if key in edge - } - mini_schema["edgelabels"].append(new_edge) - - return mini_schema - - def run(self) -> CStatus: - sts = self.node_init() - if sts.isErr(): - return sts - schema = self.schema.getSchema() - if not schema["vertexlabels"] and not schema["edgelabels"]: - raise Exception(f"Can not get {self.graph_name}'s schema from HugeGraph!") - - self.context.lock() - self.context.schema = schema - # TODO: enhance the logic here - self.context.simple_schema = self.simple_schema(schema) - self.context.unlock() - return CStatus() - - def get_result(self): - self.context.lock() - res = self.context.to_json() - self.context.unlock() - return res diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py index ee89d330f..5cdad0316 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py @@ -30,54 +30,6 @@ ) from hugegraph_llm.utils.log import log -from hugegraph_llm.operators.util import init_context -from hugegraph_llm.models.embeddings.init_embedding import get_embedding -from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState -from PyCGraph import GNode, CStatus - - -class BuildVectorIndexNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None - - def init(self): - return init_context(self) - - def node_init(self): - self.embedding = get_embedding(llm_settings) - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(self.embedding, "model_name", None) - ) - self.vector_index = VectorIndex.from_index_file( - self.index_dir, self.filename_prefix - ) - return CStatus() - - def run(self): - # init workflow input - sts = self.node_init() - if sts.isErr(): - return sts - self.context.lock() - try: - if self.context.chunks is None: - raise ValueError("chunks not found in context.") - chunks = self.context.chunks - finally: - self.context.unlock() - chunks_embedding = [] - log.debug("Building vector index for %s chunks...", len(chunks)) - # TODO: use async_get_texts_embedding instead of single sync method - chunks_embedding = asyncio.run(get_embeddings_parallel(self.embedding, chunks)) - if len(chunks_embedding) > 0: - self.vector_index.add(chunks_embedding, chunks) - self.vector_index.to_index_file(self.index_dir, self.filename_prefix) - return CStatus() - class BuildVectorIndex: def __init__(self, embedding: BaseEmbedding): 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 15a8fdda7..571ffde51 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 @@ -18,16 +18,10 @@ import re from typing import List, Any, Dict, Optional -from hugegraph_llm.config import llm_settings from hugegraph_llm.document.chunk_split import ChunkSplitter from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.utils.log import log -from hugegraph_llm.operators.util import init_context -from hugegraph_llm.models.llms.init_llm import get_chat_llm -from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState -from PyCGraph import GNode, CStatus - SCHEMA_EXAMPLE_PROMPT = """## Main Task Extract Triples from the given text and graph schema @@ -213,143 +207,3 @@ def _filter_long_id(self, graph) -> Dict[str, List[Any]]: if self.valid(edge["start"]) and self.valid(edge["end"]) ] return graph - - -class InfoExtractNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None - - def init(self): - return init_context(self) - - def node_init(self): - self.llm = get_chat_llm(llm_settings) - if self.wk_input.example_prompt is None: - return CStatus(-1, "Error occurs when prepare for workflow input") - self.example_prompt = self.wk_input.example_prompt - return CStatus() - - def extract_triples_by_regex_with_schema(self, schema, text): - text = text.replace("\\n", " ").replace("\\", " ").replace("\n", " ") - pattern = r"\((.*?), (.*?), (.*?)\) - ([^ ]*)" - matches = re.findall(pattern, text) - - vertices_dict = {v["id"]: v for v in self.context.vertices} - for match in matches: - s, p, o, label = [item.strip() for item in match] - if None in [label, s, p, o]: - continue - # TODO: use a more efficient way to compare the extract & input property - p_lower = p.lower() - for vertex in schema["vertices"]: - if vertex["vertex_label"] == label and any( - pp.lower() == p_lower for pp in vertex["properties"] - ): - id = f"{label}-{s}" - if id not in vertices_dict: - vertices_dict[id] = { - "id": id, - "name": s, - "label": label, - "properties": {p: o}, - } - else: - vertices_dict[id]["properties"].update({p: o}) - break - for edge in schema["edges"]: - if edge["edge_label"] == label: - source_label = edge["source_vertex_label"] - source_id = f"{source_label}-{s}" - if source_id not in vertices_dict: - vertices_dict[source_id] = { - "id": source_id, - "name": s, - "label": source_label, - "properties": {}, - } - target_label = edge["target_vertex_label"] - target_id = f"{target_label}-{o}" - if target_id not in vertices_dict: - vertices_dict[target_id] = { - "id": target_id, - "name": o, - "label": target_label, - "properties": {}, - } - self.context.edges.append( - { - "start": source_id, - "end": target_id, - "type": label, - "properties": {}, - } - ) - break - self.context.vertices = list(vertices_dict.values()) - - def extract_triples_by_regex(self, text): - text = text.replace("\\n", " ").replace("\\", " ").replace("\n", " ") - pattern = r"\((.*?), (.*?), (.*?)\)" - self.context.triples += re.findall(pattern, text) - - def run(self) -> CStatus: - sts = self.node_init() - if sts.isErr(): - return sts - self.context.lock() - if self.context.chunks is None: - self.context.unlock() - raise ValueError("parameter required by extract node not found in context.") - schema = self.context.schema - chunks = self.context.chunks - - if schema: - self.context.vertices = [] - self.context.edges = [] - else: - self.context.triples = [] - - self.context.unlock() - - for sentence in chunks: - proceeded_chunk = self.extract_triples_by_llm(schema, sentence) - log.debug( - "[Legacy] %s input: %s \n output:%s", - self.__class__.__name__, - sentence, - proceeded_chunk, - ) - if schema: - self.extract_triples_by_regex_with_schema(schema, proceeded_chunk) - else: - self.extract_triples_by_regex(proceeded_chunk) - - if self.context.call_count: - self.context.call_count += len(chunks) - else: - self.context.call_count = len(chunks) - self._filter_long_id() - return CStatus() - - def extract_triples_by_llm(self, schema, chunk) -> str: - prompt = generate_extract_triple_prompt(chunk, schema) - if self.example_prompt is not None: - prompt = self.example_prompt + prompt - return self.llm.generate(prompt=prompt) - - # TODO: make 'max_length' be a configurable param in settings.py/settings.cfg - def valid(self, element_id: str, max_length: int = 256) -> bool: - if len(element_id.encode("utf-8")) >= max_length: - log.warning("Filter out GraphElementID too long: %s", element_id) - return False - return True - - def _filter_long_id(self): - self.context.vertices = [ - vertex for vertex in self.context.vertices if self.valid(vertex["id"]) - ] - self.context.edges = [ - edge - for edge in self.context.edges - if self.valid(edge["start"]) and self.valid(edge["end"]) - ] 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 6e492b8f5..79fb33b4f 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 @@ -21,16 +21,11 @@ import re from typing import List, Any, Dict -from hugegraph_llm.config import llm_settings, prompt +from hugegraph_llm.config import prompt from hugegraph_llm.document.chunk_split import ChunkSplitter from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.utils.log import log -from hugegraph_llm.operators.util import init_context -from hugegraph_llm.models.llms.init_llm import get_chat_llm -from hugegraph_llm.state.ai_state import WkFlowState, WkFlowInput -from PyCGraph import GNode, CStatus - # TODO: It is not clear whether there is any other dependence on the SCHEMA_EXAMPLE_PROMPT variable. # Because the SCHEMA_EXAMPLE_PROMPT variable will no longer change based on # prompt.extract_graph_prompt changes after the system loads, this does not seem to meet expectations. @@ -182,123 +177,3 @@ def process_items(item_list, valid_labels, item_type): "Invalid property graph JSON! Please check the extracted JSON data carefully" ) return items - - -class PropertyGraphExtractNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None - - def init(self): - self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name - return init_context(self) - - def node_init(self): - self.llm = get_chat_llm(llm_settings) - if self.wk_input.example_prompt is None: - return CStatus(-1, "Error occurs when prepare for workflow input") - self.example_prompt = self.wk_input.example_prompt - return CStatus() - - def run(self) -> CStatus: - sts = self.node_init() - if sts.isErr(): - return sts - self.context.lock() - try: - if self.context.schema is None or self.context.chunks is None: - raise ValueError( - "parameter required by extract node not found in context." - ) - schema = self.context.schema - chunks = self.context.chunks - if self.context.vertices is None: - self.context.vertices = [] - if self.context.edges is None: - self.context.edges = [] - finally: - self.context.unlock() - - items = [] - for chunk in chunks: - proceeded_chunk = self.extract_property_graph_by_llm(schema, chunk) - log.debug( - "[LLM] %s input: %s \n output:%s", - self.__class__.__name__, - chunk, - proceeded_chunk, - ) - items.extend(self._extract_and_filter_label(schema, proceeded_chunk)) - items = filter_item(schema, items) - self.context.lock() - try: - for item in items: - if item["type"] == "vertex": - self.context.vertices.append(item) - elif item["type"] == "edge": - self.context.edges.append(item) - finally: - self.context.unlock() - self.context.call_count = (self.context.call_count or 0) + len(chunks) - return CStatus() - - def extract_property_graph_by_llm(self, schema, chunk): - prompt = generate_extract_property_graph_prompt(chunk, schema) - if self.example_prompt is not None: - prompt = self.example_prompt + prompt - return self.llm.generate(prompt=prompt) - - def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: - # Use regex to extract a JSON object with curly braces - json_match = re.search(r"({.*})", text, re.DOTALL) - if not json_match: - log.critical( - "Invalid property graph! No JSON object found, " - "please check the output format example in prompt." - ) - return [] - json_str = json_match.group(1).strip() - - items = [] - try: - property_graph = json.loads(json_str) - # 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'." - ) - return items - - # Create sets for valid vertex and edge labels based on the schema - vertex_label_set = {vertex["name"] for vertex in schema["vertexlabels"]} - edge_label_set = {edge["name"] for edge in schema["edgelabels"]} - - def process_items(item_list, valid_labels, item_type): - for item in item_list: - if not isinstance(item, dict): - log.warning( - "Invalid property graph item type '%s'.", type(item) - ) - continue - if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): - log.warning("Invalid item keys '%s'.", item.keys()) - continue - if item["label"] not in valid_labels: - log.warning( - "Invalid %s label '%s' has been ignored.", - item_type, - item["label"], - ) - continue - items.append(item) - - process_items(property_graph["vertices"], vertex_label_set, "vertex") - process_items(property_graph["edges"], edge_label_set, "edge") - except json.JSONDecodeError: - log.critical( - "Invalid property graph JSON! Please check the extracted JSON data carefully" - ) - return items diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 0543aa2b4..6d3418c00 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -25,6 +25,14 @@ class WkFlowInput(GParam): example_prompt: str = None # need by graph information extract schema: str = None # Schema information requeired by SchemaNode graph_name: str = None + data_json = None + extract_type = None + query_examples = None + few_shot_schema = None + # Fields related to PromptGenerate + source_text: str = None # Original text + scenario: str = None # Scenario description + example_name: str = None # Example name def reset(self, _: CStatus) -> None: self.texts = None @@ -33,6 +41,14 @@ def reset(self, _: CStatus) -> None: self.example_prompt = None self.schema = None self.graph_name = None + self.data_json = None + self.extract_type = None + self.query_examples = None + self.few_shot_schema = None + # PromptGenerate related configuration + self.source_text = None + self.scenario = None + self.example_name = None class WkFlowState(GParam): @@ -49,6 +65,8 @@ class WkFlowState(GParam): graph_result = None keywords_embeddings = None + generated_extract_prompt: Optional[str] = None + def setup(self): self.schema = None self.simple_schema = None @@ -63,6 +81,8 @@ def setup(self): self.graph_result = None self.keywords_embeddings = None + self.generated_extract_prompt = None + return CStatus() def to_json(self): @@ -79,3 +99,11 @@ def to_json(self): for k, v in self.__dict__.items() if not k.startswith("_") and v is not None } + + # Implement a method that assigns keys from data_json as WkFlowState member variables + def assign_from_json(self, data_json: dict): + """ + Assigns each key in the input json object as a member variable of WkFlowState. + """ + for k, v in data_json.items(): + setattr(self, k, v) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index f61b5f843..ccace69f2 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -36,6 +36,15 @@ def get_graph_index_info(): + try: + scheduler = SchedulerSingleton.get_instance() + return scheduler.schedule_flow("get_graph_index_info") + except Exception as e: # pylint: disable=broad-exception-caught + log.error(e) + raise gr.Error(str(e)) + + +def get_graph_index_info_old(): builder = KgBuilder( LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() ) @@ -150,6 +159,15 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: def update_vid_embedding(): + scheduler = SchedulerSingleton.get_instance() + try: + return scheduler.schedule_flow("update_vid_embeddings") + except Exception as e: # pylint: disable=broad-exception-caught + log.error(e) + raise gr.Error(str(e)) + + +def update_vid_embedding_old(): builder = KgBuilder( LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() ) @@ -166,6 +184,18 @@ def update_vid_embedding(): def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: + try: + scheduler = SchedulerSingleton.get_instance() + return scheduler.schedule_flow("import_graph_data", data, schema) + except Exception as e: # pylint: disable=W0718 + log.error(e) + traceback.print_exc() + # Note: can't use gr.Error here + gr.Warning(str(e) + " Please check the graph data format/type carefully.") + return data + + +def import_graph_data_old(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: data_json = json.loads(data.strip()) log.debug("Import graph data: %s", data) @@ -190,6 +220,16 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: def build_schema(input_text, query_example, few_shot): + scheduler = SchedulerSingleton.get_instance() + try: + return scheduler.schedule_flow( + "build_schema", input_text, query_example, few_shot + ) + except (TypeError, ValueError) as e: + raise gr.Error(f"Schema generation failed: {e}") + + +def build_schema_old(input_text, query_example, few_shot): context = { "raw_texts": [input_text] if input_text else [], "query_examples": [], From 78011d3c77f47a8919b080800dd70890c6f00720 Mon Sep 17 00:00:00 2001 From: LingXiao Qi Date: Tue, 30 Sep 2025 00:24:18 +0800 Subject: [PATCH 03/71] Refactor: text2germlin with PCgraph framework (#50) Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Linyu <94553312+weijinglin@users.noreply.github.com> --- .../src/hugegraph_llm/api/admin_api.py | 8 +- .../api/exceptions/rag_exceptions.py | 4 +- .../hugegraph_llm/api/models/rag_requests.py | 106 +++++++---- .../src/hugegraph_llm/api/rag_api.py | 43 +++-- .../src/hugegraph_llm/config/admin_config.py | 2 + .../src/hugegraph_llm/config/generate.py | 4 +- .../hugegraph_llm/config/hugegraph_config.py | 1 + .../src/hugegraph_llm/config/llm_config.py | 21 +- .../config/models/base_config.py | 22 ++- .../config/models/base_prompt_config.py | 24 ++- .../src/hugegraph_llm/config/prompt_config.py | 1 + .../demo/rag_demo/admin_block.py | 39 ++-- .../src/hugegraph_llm/demo/rag_demo/app.py | 4 +- .../demo/rag_demo/configs_block.py | 91 +++------ .../demo/rag_demo/other_block.py | 13 +- .../hugegraph_llm/demo/rag_demo/rag_block.py | 75 +++++--- .../demo/rag_demo/text2gremlin_block.py | 109 ++++++++--- .../demo/rag_demo/vector_graph_block.py | 74 +++----- .../src/hugegraph_llm/document/chunk_split.py | 14 +- .../flows/get_graph_index_info.py | 4 +- .../src/hugegraph_llm/flows/graph_extract.py | 4 +- .../hugegraph_llm/flows/import_graph_data.py | 8 +- .../hugegraph_llm/flows/prompt_generate.py | 4 +- .../src/hugegraph_llm/flows/scheduler.py | 9 +- .../src/hugegraph_llm/flows/text2gremlin.py | 112 +++++++++++ .../src/hugegraph_llm/indices/graph_index.py | 17 +- .../src/hugegraph_llm/indices/vector_index.py | 33 +++- .../hugegraph_llm/middleware/middleware.py | 5 +- .../hugegraph_llm/models/embeddings/base.py | 37 ++-- .../hugegraph_llm/models/embeddings/openai.py | 25 ++- .../src/hugegraph_llm/models/llms/base.py | 34 ++-- .../src/hugegraph_llm/models/llms/init_llm.py | 6 +- .../src/hugegraph_llm/models/llms/litellm.py | 10 +- .../src/hugegraph_llm/models/llms/ollama.py | 30 ++- .../src/hugegraph_llm/models/llms/openai.py | 4 +- .../hugegraph_llm/models/rerankers/cohere.py | 9 +- .../models/rerankers/init_reranker.py | 4 +- .../models/rerankers/siliconflow.py | 9 +- .../nodes/hugegraph_node/gremlin_execute.py | 68 +++++++ .../nodes/hugegraph_node/schema.py | 2 +- .../index_node/gremlin_example_index_query.py | 49 +++++ .../nodes/llm_node/schema_build.py | 8 +- .../nodes/llm_node/text2gremlin.py | 70 +++++++ .../operators/common_op/check_schema.py | 40 +--- .../operators/common_op/merge_dedup_rerank.py | 15 +- .../operators/document_op/word_extract.py | 3 +- .../hugegraph_llm/operators/graph_rag_task.py | 4 +- .../hugegraph_op/commit_to_hugegraph.py | 58 ++---- .../operators/hugegraph_op/graph_rag_query.py | 63 ++++-- .../operators/hugegraph_op/schema_manager.py | 4 +- .../index_op/build_gremlin_example_index.py | 14 +- .../index_op/build_semantic_index.py | 30 +-- .../operators/index_op/build_vector_index.py | 4 +- .../index_op/gremlin_example_index_query.py | 29 ++- .../operators/index_op/semantic_id_query.py | 33 ++-- .../operators/index_op/vector_index_query.py | 8 +- .../operators/kg_construction_task.py | 11 +- .../operators/llm_op/answer_synthesize.py | 179 ++++++++++++------ .../operators/llm_op/disambiguate_data.py | 3 +- .../operators/llm_op/gremlin_generate.py | 19 +- .../operators/llm_op/info_extract.py | 8 +- .../operators/llm_op/keyword_extract.py | 28 +-- .../operators/llm_op/prompt_generate.py | 6 +- .../llm_op/property_graph_extract.py | 18 +- .../operators/llm_op/schema_build.py | 14 +- .../src/hugegraph_llm/state/ai_state.py | 30 ++- .../src/hugegraph_llm/utils/anchor.py | 9 +- .../src/hugegraph_llm/utils/decorators.py | 1 + .../hugegraph_llm/utils/embedding_utils.py | 9 +- .../hugegraph_llm/utils/graph_index_utils.py | 40 +--- .../hugegraph_llm/utils/hugegraph_utils.py | 26 ++- hugegraph-llm/src/hugegraph_llm/utils/log.py | 2 +- .../hugegraph_llm/utils/vector_index_utils.py | 16 +- hugegraph-llm/src/tests/config/test_config.py | 1 + .../embeddings/test_openai_embedding.py | 1 + .../tests/models/llms/test_ollama_client.py | 7 +- .../operators/common_op/test_check_schema.py | 9 +- .../operators/common_op/test_nltk_helper.py | 1 + .../src/pyhugegraph/api/auth.py | 16 +- .../src/pyhugegraph/api/graph.py | 12 +- .../src/pyhugegraph/api/schema.py | 12 +- .../api/schema_manage/index_label.py | 12 +- .../src/pyhugegraph/api/services.py | 8 +- .../src/pyhugegraph/api/traverser.py | 45 ++--- .../src/pyhugegraph/client.py | 2 +- .../pyhugegraph/example/hugegraph_example.py | 14 +- .../structure/property_key_data.py | 4 +- .../src/pyhugegraph/utils/huge_config.py | 10 +- .../src/pyhugegraph/utils/huge_router.py | 8 +- .../src/pyhugegraph/utils/log.py | 4 +- .../src/pyhugegraph/utils/util.py | 23 ++- .../src/tests/api/test_auth.py | 8 +- .../src/tests/api/test_version.py | 8 +- .../src/tests/client_utils.py | 6 +- 94 files changed, 1322 insertions(+), 829 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/gremlin_execute.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py diff --git a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py index 05648d48e..4c192c29c 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py @@ -31,8 +31,12 @@ def admin_http_api(router: APIRouter, log_stream): @router.post("/logs", status_code=status.HTTP_200_OK) async def log_stream_api(req: LogStreamRequest): if admin_settings.admin_token != req.admin_token: - raise generate_response(RAGResponse(status_code=status.HTTP_403_FORBIDDEN, #pylint: disable=E0702 - message="Invalid admin_token")) + raise generate_response( + RAGResponse( + status_code=status.HTTP_403_FORBIDDEN, # pylint: disable=E0702 + message="Invalid admin_token", + ) + ) log_path = os.path.join("logs", req.log_file) # Create a StreamingResponse that reads from the log stream generator diff --git a/hugegraph-llm/src/hugegraph_llm/api/exceptions/rag_exceptions.py b/hugegraph-llm/src/hugegraph_llm/api/exceptions/rag_exceptions.py index 75eb14cf3..18723e30b 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/exceptions/rag_exceptions.py +++ b/hugegraph-llm/src/hugegraph_llm/api/exceptions/rag_exceptions.py @@ -21,7 +21,9 @@ class ExternalException(HTTPException): def __init__(self): - super().__init__(status_code=400, detail="Connect failed with error code -1, please check the input.") + super().__init__( + status_code=400, detail="Connect failed with error code -1, please check the input." + ) class ConnectionFailedException(HTTPException): diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py index cf227e8bd..f46aea02c 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py +++ b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py @@ -24,10 +24,10 @@ class GraphConfigRequest(BaseModel): - url: str = Query('127.0.0.1:8080', description="hugegraph client url.") - graph: str = Query('hugegraph', description="hugegraph client name.") - user: str = Query('', description="hugegraph client user.") - pwd: str = Query('', description="hugegraph client pwd.") + url: str = Query("127.0.0.1:8080", description="hugegraph client url.") + graph: str = Query("hugegraph", description="hugegraph client name.") + user: str = Query("", description="hugegraph client user.") + pwd: str = Query("", description="hugegraph client pwd.") gs: str = None @@ -36,22 +36,42 @@ class RAGRequest(BaseModel): raw_answer: bool = Query(False, description="Use LLM to generate answer directly") vector_only: bool = Query(False, description="Use LLM to generate answer with vector") graph_only: bool = Query(True, description="Use LLM to generate answer with graph RAG only") - graph_vector_answer: bool = Query(False, description="Use LLM to generate answer with vector & GraphRAG") + graph_vector_answer: bool = Query( + False, description="Use LLM to generate answer with vector & GraphRAG" + ) graph_ratio: float = Query(0.5, description="The ratio of GraphRAG ans & vector ans") - rerank_method: Literal["bleu", "reranker"] = Query("bleu", description="Method to rerank the results.") - near_neighbor_first: bool = Query(False, description="Prioritize near neighbors in the search results.") - custom_priority_info: str = Query("", description="Custom information to prioritize certain results.") + rerank_method: Literal["bleu", "reranker"] = Query( + "bleu", description="Method to rerank the results." + ) + near_neighbor_first: bool = Query( + False, description="Prioritize near neighbors in the search results." + ) + custom_priority_info: str = Query( + "", description="Custom information to prioritize certain results." + ) # Graph Configs - max_graph_items: int = Query(30, description="Maximum number of items for GQL queries in graph.") + max_graph_items: int = Query( + 30, description="Maximum number of items for GQL queries in graph." + ) topk_return_results: int = Query(20, description="Number of sorted results to return finally.") - vector_dis_threshold: float = Query(0.9, description="Threshold for vector similarity\ - (results greater than this will be ignored).") - topk_per_keyword: int = Query(1, description="TopK results returned for each keyword \ - extracted from the query, by default only the most similar one is returned.") - client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") + vector_dis_threshold: float = Query( + 0.9, + description="Threshold for vector similarity\ + (results greater than this will be ignored).", + ) + topk_per_keyword: int = Query( + 1, + description="TopK results returned for each keyword \ + extracted from the query, by default only the most similar one is returned.", + ) + client_config: Optional[GraphConfigRequest] = Query( + None, description="hugegraph server config." + ) # Keep prompt params in the end - answer_prompt: Optional[str] = Query(prompt.answer_prompt, description="Prompt to guide the answer generation.") + answer_prompt: Optional[str] = Query( + prompt.answer_prompt, description="Prompt to guide the answer generation." + ) keywords_extract_prompt: Optional[str] = Query( prompt.keywords_extract_prompt, description="Prompt for extracting keywords from query.", @@ -67,22 +87,39 @@ class RAGRequest(BaseModel): class GraphRAGRequest(BaseModel): query: str = Query(..., description="Query you want to ask") # Graph Configs - max_graph_items: int = Query(30, description="Maximum number of items for GQL queries in graph.") + max_graph_items: int = Query( + 30, description="Maximum number of items for GQL queries in graph." + ) topk_return_results: int = Query(20, description="Number of sorted results to return finally.") - vector_dis_threshold: float = Query(0.9, description="Threshold for vector similarity \ - (results greater than this will be ignored).") - topk_per_keyword: int = Query(1, description="TopK results returned for each keyword extracted\ - from the query, by default only the most similar one is returned.") + vector_dis_threshold: float = Query( + 0.9, + description="Threshold for vector similarity \ + (results greater than this will be ignored).", + ) + topk_per_keyword: int = Query( + 1, + description="TopK results returned for each keyword extracted\ + from the query, by default only the most similar one is returned.", + ) - client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") + client_config: Optional[GraphConfigRequest] = Query( + None, description="hugegraph server config." + ) get_vertex_only: bool = Query(False, description="return only keywords & vertex (early stop).") gremlin_tmpl_num: int = Query( - 1, description="Number of Gremlin templates to use. If num <=0 means template is not provided" + 1, + description="Number of Gremlin templates to use. If num <=0 means template is not provided", + ) + rerank_method: Literal["bleu", "reranker"] = Query( + "bleu", description="Method to rerank the results." + ) + near_neighbor_first: bool = Query( + False, description="Prioritize near neighbors in the search results." + ) + custom_priority_info: str = Query( + "", description="Custom information to prioritize certain results." ) - rerank_method: Literal["bleu", "reranker"] = Query("bleu", description="Method to rerank the results.") - near_neighbor_first: bool = Query(False, description="Prioritize near neighbors in the search results.") - custom_priority_info: str = Query("", description="Custom information to prioritize certain results.") gremlin_prompt: Optional[str] = Query( prompt.gremlin_generate_prompt, description="Prompt for the Text2Gremlin query.", @@ -115,6 +152,7 @@ class LogStreamRequest(BaseModel): admin_token: Optional[str] = None log_file: Optional[str] = "llm-server.log" + class GremlinOutputType(str, Enum): MATCH_RESULT = "match_result" TEMPLATE_GREMLIN = "template_gremlin" @@ -122,32 +160,36 @@ class GremlinOutputType(str, Enum): TEMPLATE_EXECUTION_RESULT = "template_execution_result" RAW_EXECUTION_RESULT = "raw_execution_result" + class GremlinGenerateRequest(BaseModel): query: str example_num: Optional[int] = Query( - 0, - description="Number of Gremlin templates to use.(0 means no templates)" + 0, description="Number of Gremlin templates to use.(0 means no templates)" ) gremlin_prompt: Optional[str] = Query( prompt.gremlin_generate_prompt, description="Prompt for the Text2Gremlin query.", ) - client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") + client_config: Optional[GraphConfigRequest] = Query( + None, description="hugegraph server config." + ) output_types: Optional[List[GremlinOutputType]] = Query( default=[GremlinOutputType.TEMPLATE_GREMLIN], description=""" a list can contain "match_result","template_gremlin", "raw_gremlin","template_execution_result","raw_execution_result" You can specify which type of result do you need. Empty means all types. - """ + """, ) - @field_validator('gremlin_prompt') + @field_validator("gremlin_prompt") @classmethod def validate_prompt_placeholders(cls, v): if v is not None: - required_placeholders = ['{query}', '{schema}', '{example}', '{vertices}'] + required_placeholders = ["{query}", "{schema}", "{example}", "{vertices}"] missing = [p for p in required_placeholders if p not in v] if missing: - raise ValueError(f"Prompt template is missing required placeholders: {', '.join(missing)}") + raise ValueError( + f"Prompt template is missing required placeholders: {', '.join(missing)}" + ) return v diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py index c39c77711..5c9295efa 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py @@ -32,6 +32,8 @@ from hugegraph_llm.api.models.rag_response import RAGResponse from hugegraph_llm.config import llm_settings, prompt from hugegraph_llm.utils.log import log +from hugegraph_llm.flows.scheduler import SchedulerSingleton + # pylint: disable=too-many-statements def rag_http_api( @@ -73,7 +75,9 @@ def rag_answer_api(req: RAGRequest): "query": req.query, **{ key: value - for key, value in zip(["raw_answer", "vector_only", "graph_only", "graph_vector_answer"], result) + for key, value in zip( + ["raw_answer", "vector_only", "graph_only", "graph_vector_answer"], result + ) if getattr(req, key) }, } @@ -102,11 +106,12 @@ def graph_rag_recall_api(req: GraphRAGRequest): near_neighbor_first=req.near_neighbor_first, custom_related_information=req.custom_priority_info, gremlin_prompt=req.gremlin_prompt or prompt.gremlin_generate_prompt, - get_vertex_only=req.get_vertex_only + get_vertex_only=req.get_vertex_only, ) if req.get_vertex_only: from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery + graph_rag = GraphRAGQuery() graph_rag.init_client(result) vertex_details = graph_rag.get_vertex_details(result["match_vids"]) @@ -134,7 +139,8 @@ def graph_rag_recall_api(req: GraphRAGRequest): except Exception as e: log.error("Unexpected error occurred: %s", e) raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="An unexpected error occurred." + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="An unexpected error occurred.", ) from e @router.post("/config/graph", status_code=status.HTTP_201_CREATED) @@ -149,7 +155,9 @@ def llm_config_api(req: LLMConfigRequest): llm_settings.llm_type = req.llm_type if req.llm_type == "openai": - res = apply_llm_conf(req.api_key, req.api_base, req.language_model, req.max_tokens, origin_call="http") + res = apply_llm_conf( + req.api_key, req.api_base, req.language_model, req.max_tokens, origin_call="http" + ) else: res = apply_llm_conf(req.host, req.port, req.language_model, None, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) @@ -159,7 +167,9 @@ def embedding_config_api(req: LLMConfigRequest): llm_settings.embedding_type = req.llm_type if req.llm_type == "openai": - res = apply_embedding_conf(req.api_key, req.api_base, req.language_model, origin_call="http") + res = apply_embedding_conf( + req.api_key, req.api_base, req.language_model, origin_call="http" + ) else: res = apply_embedding_conf(req.host, req.port, req.language_model, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) @@ -169,7 +179,9 @@ def rerank_config_api(req: RerankerConfigRequest): llm_settings.reranker_type = req.reranker_type if req.reranker_type == "cohere": - res = apply_reranker_conf(req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http") + res = apply_reranker_conf( + req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http" + ) elif req.reranker_type == "siliconflow": res = apply_reranker_conf(req.api_key, req.reranker_model, None, origin_call="http") else: @@ -181,16 +193,23 @@ def text2gremlin_api(req: GremlinGenerateRequest): try: set_graph_config(req) + # Basic parameter validation: empty query => 400 + if not req.query or not str(req.query).strip(): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail="Query must not be empty." + ) + output_types_str_list = None if req.output_types: output_types_str_list = [ot.value for ot in req.output_types] - response_dict = gremlin_generate_selective_func( - inp=req.query, - example_num=req.example_num, - schema_input=huge_settings.graph_name, - gremlin_prompt_input=req.gremlin_prompt, - requested_outputs=output_types_str_list, + response_dict = SchedulerSingleton.get_instance().schedule_flow( + "text2gremlin", + req.query, + req.example_num, + huge_settings.graph_name, + req.gremlin_prompt, + output_types_str_list, ) return response_dict except HTTPException as e: diff --git a/hugegraph-llm/src/hugegraph_llm/config/admin_config.py b/hugegraph-llm/src/hugegraph_llm/config/admin_config.py index b2814de41..fabc75de4 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/admin_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/admin_config.py @@ -18,8 +18,10 @@ from typing import Optional from .models import BaseConfig + class AdminConfig(BaseConfig): """Admin settings""" + enable_login: Optional[str] = "False" user_token: Optional[str] = "4321" admin_token: Optional[str] = "xxxx" diff --git a/hugegraph-llm/src/hugegraph_llm/config/generate.py b/hugegraph-llm/src/hugegraph_llm/config/generate.py index 36910e480..4b40e899f 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/generate.py +++ b/hugegraph-llm/src/hugegraph_llm/config/generate.py @@ -22,7 +22,9 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(description="Generate hugegraph-llm config file") - parser.add_argument("-U", "--update", default=True, action="store_true", help="Update the config file") + parser.add_argument( + "-U", "--update", default=True, action="store_true", help="Update the config file" + ) args = parser.parse_args() if args.update: huge_settings.generate_env() diff --git a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py b/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py index e51008d96..69abf0fbc 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py @@ -21,6 +21,7 @@ class HugeGraphConfig(BaseConfig): """HugeGraph settings""" + # graph server config graph_url: Optional[str] = "127.0.0.1:8080" graph_name: Optional[str] = "hugegraph" diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index b2029d983..916d70ddf 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -24,6 +24,7 @@ class LLMConfig(BaseConfig): """LLM settings""" + language: Literal["EN", "CN"] = "EN" chat_llm_type: Literal["openai", "litellm", "ollama/local"] = "openai" extract_llm_type: Literal["openai", "litellm", "ollama/local"] = "openai" @@ -31,23 +32,33 @@ class LLMConfig(BaseConfig): embedding_type: Optional[Literal["openai", "litellm", "ollama/local"]] = "openai" reranker_type: Optional[Literal["cohere", "siliconflow"]] = None # 1. OpenAI settings - openai_chat_api_base: Optional[str] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_chat_api_base: Optional[str] = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_chat_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") openai_chat_language_model: Optional[str] = "gpt-4.1-mini" - openai_extract_api_base: Optional[str] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_extract_api_base: Optional[str] = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_extract_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") openai_extract_language_model: Optional[str] = "gpt-4.1-mini" - openai_text2gql_api_base: Optional[str] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_text2gql_api_base: Optional[str] = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_text2gql_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") openai_text2gql_language_model: Optional[str] = "gpt-4.1-mini" - openai_embedding_api_base: Optional[str] = os.environ.get("OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1") + openai_embedding_api_base: Optional[str] = os.environ.get( + "OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1" + ) openai_embedding_api_key: Optional[str] = os.environ.get("OPENAI_EMBEDDING_API_KEY") openai_embedding_model: Optional[str] = "text-embedding-3-small" openai_chat_tokens: int = 8192 openai_extract_tokens: int = 256 openai_text2gql_tokens: int = 4096 # 2. Rerank settings - cohere_base_url: Optional[str] = os.environ.get("CO_API_URL", "https://api.cohere.com/v1/rerank") + cohere_base_url: Optional[str] = os.environ.get( + "CO_API_URL", "https://api.cohere.com/v1/rerank" + ) reranker_api_key: Optional[str] = None reranker_model: Optional[str] = None # 3. Ollama settings diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py index dfe9d1056..5fec3a778 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py @@ -31,12 +31,15 @@ class BaseConfig(BaseSettings): class Config: env_file = env_path case_sensitive = False - extra = 'ignore' # ignore extra fields to avoid ValidationError + extra = "ignore" # ignore extra fields to avoid ValidationError env_ignore_empty = True def generate_env(self): if os.path.exists(env_path): - log.info("%s already exists, do you want to override with the default configuration? (y/n)", env_path) + log.info( + "%s already exists, do you want to override with the default configuration? (y/n)", + env_path, + ) update = input() if update.lower() != "y": return @@ -96,8 +99,12 @@ def _sync_env_to_object(self, env_config, config_dict): obj_value_str = str(obj_value) if obj_value is not None else "" if env_value != obj_value_str: - log.info("Update configuration from the file: %s=%s (Original value: %s)", - env_key, env_value, obj_value_str) + log.info( + "Update configuration from the file: %s=%s (Original value: %s)", + env_key, + env_value, + obj_value_str, + ) # Update the object attribute (using lowercase key) setattr(self, env_key.lower(), env_value) @@ -106,8 +113,11 @@ def _sync_object_to_env(self, env_config, config_dict): for obj_key, obj_value in config_dict.items(): if obj_key not in env_config: obj_value_str = str(obj_value) if obj_value is not None else "" - log.info("Add configuration items to the environment variable file: %s=%s", - obj_key, obj_value) + log.info( + "Add configuration items to the environment variable file: %s=%s", + obj_key, + obj_value, + ) # Add to .env set_key(env_path, obj_key, obj_value_str, quote_mode="never") diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py index 7af1ef922..2369d01a6 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py @@ -32,11 +32,14 @@ class LiteralStr(str): pass + def literal_str_representer(dumper, data): - return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|') + return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|") + yaml.add_representer(LiteralStr, literal_str_representer) + class BasePromptConfig: graph_schema: str = "" extract_graph_prompt: str = "" @@ -54,9 +57,7 @@ def ensure_yaml_file_exists(self): current_dir = Path.cwd().resolve() project_root = get_project_root() if current_dir == project_root: - log.info( - "Current working directory is the project root, proceeding to run the app." - ) + log.info("Current working directory is the project root, proceeding to run the app.") else: error_msg = ( f"Current working directory is not the project root. " @@ -74,16 +75,20 @@ def ensure_yaml_file_exists(self): setattr(self, key, value) # Check if the language in the .env file matches the language in the YAML file - env_lang = (self.llm_settings.language.lower() - if hasattr(self, 'llm_settings') and self.llm_settings.language - else 'en') - yaml_lang = data.get('_language_generated', 'en').lower() + env_lang = ( + self.llm_settings.language.lower() + if hasattr(self, "llm_settings") and self.llm_settings.language + else "en" + ) + yaml_lang = data.get("_language_generated", "en").lower() if env_lang.strip() != yaml_lang.strip(): log.warning( "Prompt was changed '.env' language is '%s', " "but '%s' was generated for '%s'. " "Regenerating the prompt file...", - env_lang, F_NAME, yaml_lang + env_lang, + F_NAME, + yaml_lang, ) if self.llm_settings.language.lower() == "cn": self.answer_prompt = self.answer_prompt_CN @@ -105,6 +110,7 @@ def save_to_yaml(self): def to_literal(val): return LiteralStr(val) if isinstance(val, str) else val + data = { "graph_schema": to_literal(self.graph_schema), "text2gql_graph_schema": to_literal(self.text2gql_graph_schema), diff --git a/hugegraph-llm/src/hugegraph_llm/config/prompt_config.py b/hugegraph-llm/src/hugegraph_llm/config/prompt_config.py index e5e1c9267..7105f24b1 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/prompt_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/prompt_config.py @@ -23,6 +23,7 @@ class PromptConfig(BasePromptConfig): def __init__(self, llm_config_object): self.llm_settings = llm_config_object + # Data is detached from llm_op/answer_synthesize.py answer_prompt_EN: str = """You are an expert in the fields of knowledge graphs and natural language processing. diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 2d5937a43..1b2032b23 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -30,7 +30,7 @@ async def log_stream(log_path: str, lines: int = 125): Stream the content of a log file like `tail -f`. """ try: - with open(log_path, 'r', encoding='utf-8') as file: + with open(log_path, "r", encoding="utf-8") as file: buffer = deque(file, maxlen=lines) for line in buffer: yield line # Yield the initial lines @@ -50,8 +50,8 @@ async def log_stream(log_path: str, lines: int = 125): def read_llm_server_log(lines=250): log_path = "logs/llm-server.log" try: - with open(log_path, "r", encoding='utf-8', errors="replace") as f: - return ''.join(deque(f, maxlen=lines)) + with open(log_path, "r", encoding="utf-8", errors="replace") as f: + return "".join(deque(f, maxlen=lines)) except FileNotFoundError: log.critical("Log file not found: %s", log_path) return "LLM Server log file not found." @@ -61,10 +61,10 @@ def read_llm_server_log(lines=250): def clear_llm_server_log(): log_path = "logs/llm-server.log" try: - with open(log_path, "w", encoding='utf-8') as f: + with open(log_path, "w", encoding="utf-8") as f: f.truncate(0) # Clear the contents of the file return "LLM Server log cleared." - except Exception as e: #pylint: disable=W0718 + except Exception as e: # pylint: disable=W0718 log.error("An error occurred while clearing the log: %s", str(e)) return "Failed to clear LLM Server log." @@ -84,7 +84,7 @@ def check_password(password, request: Request = None): gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), - gr.update(visible=False) + gr.update(visible=False), ) # Log the failed attempt with IP address log.error("Incorrect password attempt from IP: %s", client_ip) @@ -93,7 +93,7 @@ def check_password(password, request: Request = None): gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), - gr.update(value="Incorrect password. Access denied.", visible=True) + gr.update(value="Incorrect password. Access denied.", visible=True), ) @@ -110,10 +110,7 @@ def create_admin_block(): # Error message box, initially hidden error_message = gr.Textbox( - label="", - visible=False, - interactive=False, - elem_classes="error-message" + label="", visible=False, interactive=False, elem_classes="error-message" ) # Button to submit password @@ -136,26 +133,32 @@ def create_admin_block(): clear_llm_server_button = gr.Button("Clear LLM Server Log", visible=False) with gr.Column(): # Button to refresh LLM Server log manually - refresh_llm_server_button = gr.Button("Refresh LLM Server Log", visible=False, - variant="primary") + refresh_llm_server_button = gr.Button( + "Refresh LLM Server Log", visible=False, variant="primary" + ) # Define what happens when the password is submitted - submit_button.click( #pylint: disable=E1101 + submit_button.click( # pylint: disable=E1101 fn=check_password, inputs=[password_input], - outputs=[llm_server_log_output, hidden_row, clear_llm_server_button, - refresh_llm_server_button, error_message], + outputs=[ + llm_server_log_output, + hidden_row, + clear_llm_server_button, + refresh_llm_server_button, + error_message, + ], ) # Define what happens when the Clear LLM Server Log button is clicked - clear_llm_server_button.click( #pylint: disable=E1101 + clear_llm_server_button.click( # pylint: disable=E1101 fn=clear_llm_server_log, inputs=[], outputs=[llm_server_log_output], ) # Define what happens when the Refresh LLM Server Log button is clicked - refresh_llm_server_button.click( #pylint: disable=E1101 + refresh_llm_server_button.click( # pylint: disable=E1101 fn=read_llm_server_log, inputs=[], outputs=[llm_server_log_output], 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 2f9c3b345..eabd5dd9b 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -165,9 +165,7 @@ def create_app(): # settings.check_env() prompt.update_yaml_file() auth_enabled = admin_settings.enable_login.lower() == "true" - log.info( - "(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled" - ) + log.info("(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled") api_auth = APIRouter(dependencies=[Depends(authenticate)] if auth_enabled else []) hugegraph_llm = init_rag_ui() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 01ea24aa8..8c595c30d 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -71,9 +71,7 @@ def test_api_connection( log.debug("Request URL: %s", url) try: if method.upper() == "GET": - resp = requests.get( - url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth - ) + resp = requests.get(url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth) elif method.upper() == "POST": resp = requests.post( url, @@ -125,9 +123,7 @@ def apply_embedding_config(arg1, arg2, arg3, origin_call=None) -> int: llm_settings.ollama_embedding_host = arg1 llm_settings.ollama_embedding_port = int(arg2) llm_settings.ollama_embedding_model = arg3 - status_code = test_api_connection( - f"http://{arg1}:{arg2}", origin_call=origin_call - ) + status_code = test_api_connection(f"http://{arg1}:{arg2}", origin_call=origin_call) elif embedding_option == "litellm": llm_settings.litellm_embedding_api_key = arg1 llm_settings.litellm_embedding_api_base = arg2 @@ -218,8 +214,7 @@ def apply_llm_config( setattr(llm_settings, f"openai_{current_llm_config}_tokens", int(max_tokens)) test_url = ( - getattr(llm_settings, f"openai_{current_llm_config}_api_base") - + "/chat/completions" + getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" ) data = { "model": model_name, @@ -233,9 +228,7 @@ def apply_llm_config( elif llm_option == "ollama/local": setattr(llm_settings, f"ollama_{current_llm_config}_host", api_key_or_host) - setattr( - llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port) - ) + setattr(llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port)) setattr(llm_settings, f"ollama_{current_llm_config}_language_model", model_name) status_code = test_api_connection( f"http://{api_key_or_host}:{api_base_or_port}", origin_call=origin_call @@ -243,12 +236,8 @@ def apply_llm_config( elif llm_option == "litellm": setattr(llm_settings, f"litellm_{current_llm_config}_api_key", api_key_or_host) - setattr( - llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port - ) - setattr( - llm_settings, f"litellm_{current_llm_config}_language_model", model_name - ) + setattr(llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port) + setattr(llm_settings, f"litellm_{current_llm_config}_language_model", model_name) setattr(llm_settings, f"litellm_{current_llm_config}_tokens", int(max_tokens)) status_code = test_litellm_chat( @@ -295,7 +284,9 @@ def create_configs_block() -> list: ), ] graph_config_button = gr.Button("Apply Configuration") - graph_config_button.click(apply_graph_config, inputs=graph_config_input) # pylint: disable=no-member + graph_config_button.click( + apply_graph_config, inputs=graph_config_input + ) # pylint: disable=no-member # TODO : use OOP to refactor the following code with gr.Accordion("2. Set up the LLM.", open=False): @@ -373,13 +364,9 @@ def chat_llm_settings(llm_type): ), ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_chat_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_chat_op, inputs=llm_config_input) # Determine whether there are Settings in the.env file env_path = os.path.join( os.getcwd(), ".env" @@ -419,9 +406,7 @@ def extract_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr( - llm_settings, "openai_extract_language_model" - ), + value=getattr(llm_settings, "openai_extract_language_model"), label="model_name", ), gr.Textbox( @@ -440,9 +425,7 @@ def extract_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr( - llm_settings, "ollama_extract_language_model" - ), + value=getattr(llm_settings, "ollama_extract_language_model"), label="model_name", ), gr.Textbox(value="", visible=False), @@ -460,9 +443,7 @@ def extract_llm_settings(llm_type): info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr( - llm_settings, "litellm_extract_language_model" - ), + value=getattr(llm_settings, "litellm_extract_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -472,13 +453,9 @@ def extract_llm_settings(llm_type): ), ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_extract_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) with gr.Tab(label="text2gql"): text2gql_llm_dropdown = gr.Dropdown( @@ -503,9 +480,7 @@ def text2gql_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr( - llm_settings, "openai_text2gql_language_model" - ), + value=getattr(llm_settings, "openai_text2gql_language_model"), label="model_name", ), gr.Textbox( @@ -524,9 +499,7 @@ def text2gql_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr( - llm_settings, "ollama_text2gql_language_model" - ), + value=getattr(llm_settings, "ollama_text2gql_language_model"), label="model_name", ), gr.Textbox(value="", visible=False), @@ -544,9 +517,7 @@ def text2gql_llm_settings(llm_type): info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr( - llm_settings, "litellm_text2gql_language_model" - ), + value=getattr(llm_settings, "litellm_text2gql_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -556,13 +527,9 @@ def text2gql_llm_settings(llm_type): ), ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_text2gql_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_text2gql_op, inputs=llm_config_input) with gr.Accordion("3. Set up the Embedding.", open=False): embedding_dropdown = gr.Dropdown( @@ -594,12 +561,8 @@ def embedding_settings(embedding_type): elif embedding_type == "ollama/local": with gr.Row(): embedding_config_input = [ - gr.Textbox( - value=llm_settings.ollama_embedding_host, label="host" - ), - gr.Textbox( - value=str(llm_settings.ollama_embedding_port), label="port" - ), + gr.Textbox(value=llm_settings.ollama_embedding_host, label="host"), + gr.Textbox(value=str(llm_settings.ollama_embedding_port), label="port"), gr.Textbox( value=llm_settings.ollama_embedding_model, label="model_name", @@ -648,9 +611,7 @@ def embedding_settings(embedding_type): @gr.render(inputs=[reranker_dropdown]) def reranker_settings(reranker_type): - llm_settings.reranker_type = ( - reranker_type if reranker_type != "None" else None - ) + llm_settings.reranker_type = reranker_type if reranker_type != "None" else None if reranker_type == "cohere": with gr.Row(): reranker_config_input = [ @@ -660,9 +621,7 @@ def reranker_settings(reranker_type): type="password", ), gr.Textbox(value=llm_settings.reranker_model, label="model"), - gr.Textbox( - value=llm_settings.cohere_base_url, label="base_url" - ), + gr.Textbox(value=llm_settings.cohere_base_url, label="base_url"), ] elif reranker_type == "siliconflow": with gr.Row(): diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py index da10f50f4..8b78328f3 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py @@ -31,7 +31,9 @@ def create_other_block(): gr.Markdown("""## Other Tools """) with gr.Row(): - inp = gr.Textbox(value="g.V().limit(10)", label="Gremlin query", show_copy_button=True, lines=8) + inp = gr.Textbox( + value="g.V().limit(10)", label="Gremlin query", show_copy_button=True, lines=8 + ) out = gr.Code(label="Output", language="json", elem_classes="code-container-show") btn = gr.Button("Run Gremlin query") btn.click(fn=run_gremlin_query, inputs=[inp], outputs=out) # pylint: disable=no-member @@ -39,7 +41,9 @@ def create_other_block(): gr.Markdown("---") with gr.Row(): inp = [] - out = gr.Textbox(label="Backup Graph Manually (Auto backup at 1:00 AM everyday)", show_copy_button=True) + out = gr.Textbox( + label="Backup Graph Manually (Auto backup at 1:00 AM everyday)", show_copy_button=True + ) btn = gr.Button("Backup Graph Data") btn.click(fn=backup_data, inputs=inp, outputs=out) # pylint: disable=no-member with gr.Accordion("Init HugeGraph test data (🚧)", open=False): @@ -55,10 +59,7 @@ async def lifespan(app: FastAPI): # pylint: disable=W0621 log.info("Starting background scheduler...") scheduler = AsyncIOScheduler() scheduler.add_job( - backup_data, - trigger=CronTrigger(hour=1, minute=0), - id="daily_backup", - replace_existing=True + backup_data, trigger=CronTrigger(hour=1, minute=0), id="daily_backup", replace_existing=True ) scheduler.start() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index cc6bb44ea..c93ec5739 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -30,6 +30,7 @@ from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize from hugegraph_llm.utils.log import log + def rag_answer( text: str, raw_answer: bool, @@ -90,7 +91,9 @@ def rag_answer( near_neighbor_first=near_neighbor_first, topk_return_results=topk_return_results, ) - rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) + rag.synthesize_answer( + raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt + ) try: context = rag.run( @@ -145,6 +148,7 @@ def update_ui_configs( graph_search = graph_only_answer or graph_vector_answer return graph_search, gremlin_prompt, vector_search + async def rag_answer_streaming( text: str, raw_answer: bool, @@ -187,9 +191,9 @@ async def rag_answer_streaming( if vector_search: rag.query_vector_index() if graph_search: - rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid().import_schema( - huge_settings.graph_name - ).query_graphdb( + rag.extract_keywords( + extract_template=keywords_extract_prompt + ).keywords_to_vid().import_schema(huge_settings.graph_name).query_graphdb( num_gremlin_generate_example=gremlin_tmpl_num, gremlin_prompt=gremlin_prompt, ) @@ -201,7 +205,9 @@ async def rag_answer_streaming( # rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) try: - context = rag.run(verbose=True, query=text, vector_search=vector_search, graph_search=graph_search) + context = rag.run( + verbose=True, query=text, vector_search=vector_search, graph_search=graph_search + ) if context.get("switch_to_bleu"): gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") answer_synthesize = AnswerSynthesize( @@ -227,6 +233,7 @@ async def rag_answer_streaming( log.critical(e) raise gr.Error(f"An unexpected error occurred: {str(e)}") + @with_task_id def create_rag_block(): # pylint: disable=R0915 (too-many-statements),C0301 @@ -234,7 +241,9 @@ def create_rag_block(): with gr.Row(): with gr.Column(scale=2): # with gr.Blocks().queue(max_size=20, default_concurrency_limit=5): - inp = gr.Textbox(value=prompt.default_question, label="Question", show_copy_button=True, lines=3) + inp = gr.Textbox( + value=prompt.default_question, label="Question", show_copy_button=True, lines=3 + ) # TODO: Only support inline formula now. Should support block formula gr.Markdown("Basic LLM Answer", elem_classes="output-box-label") @@ -275,10 +284,16 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") - vector_only_radio = gr.Radio(choices=[True, False], value=False, label="Vector-only Answer") + vector_only_radio = gr.Radio( + choices=[True, False], value=False, label="Vector-only Answer" + ) with gr.Row(): - graph_only_radio = gr.Radio(choices=[True, False], value=True, label="Graph-only Answer") - graph_vector_radio = gr.Radio(choices=[True, False], value=False, label="Graph-Vector Answer") + graph_only_radio = gr.Radio( + choices=[True, False], value=True, label="Graph-only Answer" + ) + graph_vector_radio = gr.Radio( + choices=[True, False], value=False, label="Graph-Vector Answer" + ) def toggle_slider(enable): return gr.update(interactive=enable) @@ -291,8 +306,12 @@ def toggle_slider(enable): value="reranker" if online_rerank else "bleu", label="Rerank method", ) - example_num = gr.Number(value=-1, label="Template Num (<0 means disable text2gql) ", precision=0) - graph_ratio = gr.Slider(0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False) + example_num = gr.Number( + value=-1, label="Template Num (<0 means disable text2gql) ", precision=0 + ) + graph_ratio = gr.Slider( + 0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False + ) graph_vector_radio.change( toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio @@ -325,8 +344,8 @@ def toggle_slider(enable): example_num, ], outputs=[raw_out, vector_only_out, graph_only_out, graph_vector_out], - queue=True, # Enable queueing for this event - concurrency_limit=5, # Maximum of 5 concurrent executions + queue=True, # Enable queueing for this event + concurrency_limit=5, # Maximum of 5 concurrent executions ) gr.Markdown( @@ -394,18 +413,20 @@ def several_rag_answer( total_rows = len(df) for index, row in df.iterrows(): question = row.iloc[0] - basic_llm_answer, vector_only_answer, graph_only_answer, graph_vector_answer = rag_answer( - question, - is_raw_answer, - is_vector_only_answer, - is_graph_only_answer, - is_graph_vector_answer, - graph_ratio_ui, - rerank_method_ui, - near_neighbor_first_ui, - custom_related_information_ui, - answer_prompt, - keywords_extract_prompt, + basic_llm_answer, vector_only_answer, graph_only_answer, graph_vector_answer = ( + rag_answer( + question, + is_raw_answer, + is_vector_only_answer, + is_graph_only_answer, + is_graph_vector_answer, + graph_ratio_ui, + rerank_method_ui, + near_neighbor_first_ui, + custom_related_information_ui, + answer_prompt, + keywords_extract_prompt, + ) ) df.at[index, "Basic LLM Answer"] = basic_llm_answer df.at[index, "Vector-only Answer"] = vector_only_answer @@ -418,7 +439,9 @@ def several_rag_answer( with gr.Row(): with gr.Column(): - questions_file = gr.File(file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)") + questions_file = gr.File( + file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" + ) with gr.Column(): test_template_file = os.path.join(resource_path, "demo", "questions_template.xlsx") gr.File(value=test_template_file, label="Download Template File") diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 7d682403f..6600d7c41 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -33,11 +33,13 @@ from hugegraph_llm.utils.embedding_utils import get_index_folder_name from hugegraph_llm.utils.hugegraph_utils import run_gremlin_query from hugegraph_llm.utils.log import log +from hugegraph_llm.flows.scheduler import SchedulerSingleton @dataclass class GremlinResult: """Standardized result class for gremlin_generate function""" + success: bool match_result: str template_gremlin: Optional[str] = None @@ -47,13 +49,19 @@ class GremlinResult: error_message: Optional[str] = None @classmethod - def error(cls, message: str) -> 'GremlinResult': + def error(cls, message: str) -> "GremlinResult": """Create an error result""" return cls(success=False, match_result=message, error_message=message) @classmethod - def success_result(cls, match_result: str, template_gremlin: str, - raw_gremlin: str, template_exec: str, raw_exec: str) -> 'GremlinResult': + def success_result( + cls, + match_result: str, + template_gremlin: str, + raw_gremlin: str, + template_exec: str, + raw_exec: str, + ) -> "GremlinResult": """Create a successful result""" return cls( success=True, @@ -61,7 +69,7 @@ def success_result(cls, match_result: str, template_gremlin: str, template_gremlin=template_gremlin, raw_gremlin=raw_gremlin, template_exec_result=template_exec, - raw_exec_result=raw_exec + raw_exec_result=raw_exec, ) @@ -93,6 +101,7 @@ def build_example_vector_index(temp_file) -> dict: target_file = os.path.join(resource_path, folder_name, "gremlin_examples", file_name) try: import shutil + shutil.copy2(full_path, target_file) log.info("Successfully copied file to: %s", target_file) except (OSError, IOError) as e: @@ -143,7 +152,7 @@ def _configure_output_types(requested_outputs): "template_gremlin": True, "raw_gremlin": True, "template_execution_result": True, - "raw_execution_result": True + "raw_execution_result": True, } if requested_outputs: for key in output_types: @@ -176,7 +185,9 @@ def _execute_queries(context, output_types): def gremlin_generate( inp, example_num, schema, gremlin_prompt, requested_outputs: Optional[List[str]] = None ) -> GremlinResult: - generator = GremlinGenerator(llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding()) + generator = GremlinGenerator( + llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() + ) sm = SchemaManager(graph_name=schema) processed_schema, short_schema = _process_schema(schema, generator, sm) @@ -196,7 +207,9 @@ def gremlin_generate( _execute_queries(context, output_types) - match_result = json.dumps(context.get("match_result", "No Results"), ensure_ascii=False, indent=2) + match_result = json.dumps( + context.get("match_result", "No Results"), ensure_ascii=False, indent=2 + ) return GremlinResult.success_result( match_result=match_result, template_gremlin=context["result"], @@ -220,7 +233,11 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: if "edgelabels" in schema: mini_schema["edgelabels"] = [] for edge in schema["edgelabels"]: - new_edge = {key: edge[key] for key in ["name", "source_label", "target_label", "properties"] if key in edge} + new_edge = { + key: edge[key] + for key in ["name", "source_label", "target_label", "properties"] + if key in edge + } mini_schema["edgelabels"].append(new_edge) return mini_schema @@ -228,17 +245,40 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: def gremlin_generate_for_ui(inp, example_num, schema, gremlin_prompt): """UI wrapper for gremlin_generate that returns tuple for Gradio compatibility""" - result = gremlin_generate(inp, example_num, schema, gremlin_prompt) - - if not result.success: - return result.match_result, "", "", "", "" + # Execute via scheduler + try: + res = SchedulerSingleton.get_instance().schedule_flow( + "text2gremlin", + inp, + int(example_num) if isinstance(example_num, (int, float, str)) else 2, + schema, + gremlin_prompt, + [ + "match_result", + "template_gremlin", + "raw_gremlin", + "template_execution_result", + "raw_execution_result", + ], + ) + except Exception as e: # pylint: disable=broad-except + log.error("UI text2gremlin error: %s", e) + return json.dumps({"error": str(e)}, ensure_ascii=False), "", "", "", "" + + # Backward-compatible mapping for outputs + match_result = res.get("match_result", []) + match_result_str = ( + json.dumps(match_result, ensure_ascii=False, indent=2) + if isinstance(match_result, (list, dict)) + else str(match_result) + ) return ( - result.match_result, - result.template_gremlin or "", - result.raw_gremlin or "", - result.template_exec_result or "", - result.raw_exec_result or "" + match_result_str, + res.get("template_gremlin", "") or "", + res.get("raw_gremlin", "") or "", + res.get("template_execution_result", "") or "", + res.get("raw_execution_result", "") or "", ) @@ -253,7 +293,8 @@ def create_text2gremlin_block() -> Tuple: ) with gr.Row(): file = gr.File( - value=os.path.join(resource_path, "demo", "text2gremlin.csv"), label="Upload Text-Gremlin Pairs File" + value=os.path.join(resource_path, "demo", "text2gremlin.csv"), + label="Upload Text-Gremlin Pairs File", ) out = gr.Textbox(label="Result Message") with gr.Row(): @@ -263,22 +304,39 @@ def create_text2gremlin_block() -> Tuple: with gr.Row(): with gr.Column(scale=1): - input_box = gr.Textbox(value=prompt.default_question, label="Nature Language Query", show_copy_button=True) - match = gr.Code(label="Similar Template (TopN)", language="javascript", elem_classes="code-container-show") + input_box = gr.Textbox( + value=prompt.default_question, label="Nature Language Query", show_copy_button=True + ) + match = gr.Code( + label="Similar Template (TopN)", + language="javascript", + elem_classes="code-container-show", + ) initialized_out = gr.Textbox(label="Gremlin With Template", show_copy_button=True) raw_out = gr.Textbox(label="Gremlin Without Template", show_copy_button=True) tmpl_exec_out = gr.Code( - label="Query With Template Output", language="json", elem_classes="code-container-show" + label="Query With Template Output", + language="json", + elem_classes="code-container-show", ) raw_exec_out = gr.Code( - label="Query Without Template Output", language="json", elem_classes="code-container-show" + label="Query Without Template Output", + language="json", + elem_classes="code-container-show", ) with gr.Column(scale=1): - example_num_slider = gr.Slider(minimum=0, maximum=10, step=1, value=2, label="Number of refer examples") - schema_box = gr.Textbox(value=prompt.text2gql_graph_schema, label="Schema", lines=2, show_copy_button=True) + example_num_slider = gr.Slider( + minimum=0, maximum=10, step=1, value=2, label="Number of refer examples" + ) + schema_box = gr.Textbox( + value=prompt.text2gql_graph_schema, label="Schema", lines=2, show_copy_button=True + ) prompt_box = gr.Textbox( - value=prompt.gremlin_generate_prompt, label="Prompt", lines=20, show_copy_button=True + value=prompt.gremlin_generate_prompt, + label="Prompt", + lines=20, + show_copy_button=True, ) btn = gr.Button("Text2Gremlin", variant="primary") btn.click( # pylint: disable=no-member @@ -324,6 +382,7 @@ def graph_rag_recall( context = rag.run(verbose=True, query=query, graph_search=True) return context + def gremlin_generate_selective( inp: str, example_num: int, diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 4aa476942..56b5de4b3 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -63,16 +63,12 @@ def generate_prompt_for_ui(source_text, scenario, example_name): Handles the UI logic for generating a new prompt using the new workflow architecture. """ if not all([source_text, scenario, example_name]): - gr.Warning( - "Please provide original text, expected scenario, and select an example!" - ) + gr.Warning("Please provide original text, expected scenario, and select an example!") return gr.update() try: # using new architecture scheduler = SchedulerSingleton.get_instance() - result = scheduler.schedule_flow( - "prompt_generate", source_text, scenario, example_name - ) + result = scheduler.schedule_flow("prompt_generate", source_text, scenario, example_name) gr.Info("Prompt generated successfully!") return result except Exception as e: @@ -83,9 +79,7 @@ def generate_prompt_for_ui(source_text, scenario, example_name): def load_example_names(): """Load all candidate examples""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "prompt_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return [example.get("name", "Unnamed example") for example in examples] @@ -99,27 +93,23 @@ def load_query_examples(): language = getattr( prompt, "language", - getattr(prompt.llm_settings, "language", "EN") - if hasattr(prompt, "llm_settings") - else "EN", + ( + getattr(prompt.llm_settings, "language", "EN") + if hasattr(prompt, "llm_settings") + else "EN" + ), ) if language.upper() == "CN": - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples_CN.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples_CN.json") else: - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) except (FileNotFoundError, json.JSONDecodeError): try: - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -130,9 +120,7 @@ def load_query_examples(): def load_schema_fewshot_examples(): """Load few-shot examples from a JSON file""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "schema_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "schema_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -143,14 +131,10 @@ def load_schema_fewshot_examples(): def update_example_preview(example_name): """Update the display content based on the selected example name.""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "prompt_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") with open(examples_path, "r", encoding="utf-8") as f: all_examples = json.load(f) - selected_example = next( - (ex for ex in all_examples if ex.get("name") == example_name), None - ) + selected_example = next((ex for ex in all_examples if ex.get("name") == example_name), None) if selected_example: return ( @@ -178,9 +162,11 @@ def _create_prompt_helper_block(demo, input_text, info_extract_template): few_shot_dropdown = gr.Dropdown( choices=example_names, label="Select a Few-shot example as a reference", - value=example_names[0] - if example_names and example_names[0] != "No available examples" - else None, + value=( + example_names[0] + if example_names and example_names[0] != "No available examples" + else None + ), ) with gr.Accordion("View example details", open=False): example_desc_preview = gr.Markdown(label="Example description") @@ -193,9 +179,7 @@ def _create_prompt_helper_block(demo, input_text, info_extract_template): interactive=False, ) - generate_prompt_btn = gr.Button( - "🚀 Auto-generate Graph Extract Prompt", variant="primary" - ) + generate_prompt_btn = gr.Button("🚀 Auto-generate Graph Extract Prompt", variant="primary") # Bind the change event of the dropdown menu few_shot_dropdown.change( fn=update_example_preview, @@ -287,9 +271,7 @@ def create_vector_graph_block(): lines=15, max_lines=29, ) - out = gr.Code( - label="Output Info", language="json", elem_classes="code-container-edit" - ) + out = gr.Code(label="Output Info", language="json", elem_classes="code-container-edit") with gr.Row(): with gr.Accordion("Get RAG Info", open=False): @@ -298,12 +280,8 @@ def create_vector_graph_block(): graph_index_btn0 = gr.Button("Get Graph Index Info", size="sm") with gr.Accordion("Clear RAG Data", open=False): with gr.Column(): - vector_index_btn1 = gr.Button( - "Clear Chunks Vector Index", size="sm" - ) - graph_index_btn1 = gr.Button( - "Clear Graph Vid Vector Index", size="sm" - ) + vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") + graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") vector_import_bt = gr.Button("Import into Vector", variant="primary") @@ -376,9 +354,9 @@ def create_vector_graph_block(): inputs=[input_text, input_schema, info_extract_template], ) - graph_loading_bt.click( - import_graph_data, inputs=[out, input_schema], outputs=[out] - ).then(update_vid_embedding).then( + graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( + update_vid_embedding + ).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) diff --git a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py index 495ef667c..ee173b284 100644 --- a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py +++ b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py @@ -22,9 +22,9 @@ class ChunkSplitter: def __init__( - self, - split_type: Literal["paragraph", "sentence"] = "paragraph", - language: Literal["zh", "en"] = "zh" + self, + split_type: Literal["paragraph", "sentence"] = "paragraph", + language: Literal["zh", "en"] = "zh", ): if language == "zh": separators = ["\n\n", "\n", "。", "，", ""] @@ -34,15 +34,11 @@ def __init__( raise ValueError("Argument `language` must be zh or en!") if split_type == "paragraph": self.text_splitter = RecursiveCharacterTextSplitter( - chunk_size=500, - chunk_overlap=30, - separators=separators + chunk_size=500, chunk_overlap=30, separators=separators ) elif split_type == "sentence": self.text_splitter = RecursiveCharacterTextSplitter( - chunk_size=50, - chunk_overlap=0, - separators=separators + chunk_size=50, chunk_overlap=0, separators=separators ) else: raise ValueError("Arg `type` must be paragraph, sentence!") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py index fa10d0199..7d2735352 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py @@ -48,9 +48,7 @@ def build_flow(self, *args, **kwargs): def post_deal(self, pipeline=None): graph_summary_info = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) filename_prefix = get_filename_prefix( llm_settings.embedding_type, diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py index 1b0c98253..55f53b7ad 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py @@ -27,9 +27,7 @@ class GraphExtractFlow(BaseFlow): def __init__(self): pass - def prepare( - self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type - ): + def prepare(self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type): # prepare input data prepared_input.texts = texts prepared_input.language = "zh" 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 5581ef107..0b29b4e64 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py @@ -35,9 +35,11 @@ def prepare(self, prepared_input: WkFlowInput, data, schema): raise ValueError(f"Invalid JSON for 'data': {e.msg}") from e log.debug( "Import graph data (truncated): %s", - (data[:512] + "...") - if isinstance(data, str) and len(data) > 512 - else (data if isinstance(data, str) else ""), + ( + (data[:512] + "...") + if isinstance(data, str) and len(data) > 512 + else (data if isinstance(data, str) else "") + ), ) prepared_input.data_json = data_json prepared_input.schema = schema diff --git a/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py index aece6bd61..b4a7bf329 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py @@ -58,6 +58,4 @@ def post_deal(self, pipeline=None): Process the execution result of PromptGenerate workflow """ res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - return res.get( - "generated_extract_prompt", "Generation failed. Please check the logs." - ) + return res.get("generated_extract_prompt", "Generation failed. Please check the logs.") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index 559540ce3..3aedbe7f2 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -25,6 +25,7 @@ from hugegraph_llm.flows.build_schema import BuildSchemaFlow from hugegraph_llm.flows.prompt_generate import PromptGenerateFlow from hugegraph_llm.utils.log import log +from hugegraph_llm.flows.text2gremlin import Text2GremlinFlow class Scheduler: @@ -62,6 +63,10 @@ def __init__(self, max_pipeline: int = 10): "manager": GPipelineManager(), "flow": PromptGenerateFlow(), } + self.pipeline_pool["text2gremlin"] = { + "manager": GPipelineManager(), + "flow": Text2GremlinFlow(), + } self.max_pipeline = max_pipeline # TODO: Implement Agentic Workflow @@ -96,7 +101,9 @@ def schedule_flow(self, flow: str, *args, **kwargs): flow.prepare(prepared_input, *args, **kwargs) status = pipeline.run() if status.isErr(): - raise RuntimeError(f"Error in flow execution {status.getInfo()}") + error_msg = f"Error in flow execution {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) res = flow.post_deal(pipeline) manager.release(pipeline) return res diff --git a/hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py new file mode 100644 index 000000000..e9ba4276c --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py @@ -0,0 +1,112 @@ +# 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 PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.nodes.index_node.gremlin_example_index_query import GremlinExampleIndexQueryNode +from hugegraph_llm.nodes.llm_node.text2gremlin import Text2GremlinNode +from hugegraph_llm.nodes.hugegraph_node.gremlin_execute import GremlinExecuteNode + +from typing import Any, Dict, List, Optional + + +class Text2GremlinFlow(BaseFlow): + def __init__(self): + pass + + def prepare( + self, + prepared_input: WkFlowInput, + query: str, + example_num: int, + schema_input: str, + gremlin_prompt_input: Optional[str], + requested_outputs: Optional[List[str]], + ): + # sanitize example_num to [0,10], fallback to 2 if invalid + if not isinstance(example_num, int): + example_num = 2 + example_num = max(0, min(10, example_num)) + + # filter requested_outputs to allowed set and cap to 5 + allowed = { + "match_result", + "template_gremlin", + "raw_gremlin", + "template_execution_result", + "raw_execution_result", + } + req = requested_outputs or ["template_gremlin"] + req = [x for x in req if x in allowed] + if not req: + req = ["template_gremlin"] + if len(req) > 5: + req = req[:5] + + prepared_input.query = query + prepared_input.example_num = example_num + prepared_input.schema = schema_input + prepared_input.gremlin_prompt = gremlin_prompt_input + prepared_input.requested_outputs = req + return + + def build_flow( + self, + query: str, + example_num: int, + schema_input: str, + gremlin_prompt_input: Optional[str] = None, + requested_outputs: Optional[List[str]] = None, + ): + pipeline = GPipeline() + + prepared_input = WkFlowInput() + self.prepare( + prepared_input, + query=query, + example_num=example_num, + schema_input=schema_input, + gremlin_prompt_input=gremlin_prompt_input, + requested_outputs=requested_outputs, + ) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + schema_node = SchemaNode() + ieq_node = GremlinExampleIndexQueryNode() + tgn_node = Text2GremlinNode() + exe_node = GremlinExecuteNode() + + pipeline.registerGElement(schema_node, set(), "schema_node") + pipeline.registerGElement(ieq_node, set(), "gremlin_example_index_query") + pipeline.registerGElement(tgn_node, {schema_node, ieq_node}, "text2gremlin") + pipeline.registerGElement(exe_node, {tgn_node}, "gremlin_execute") + + return pipeline + + def post_deal(self, pipeline=None) -> Dict[str, Any]: + state = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + # 始终返回 5 个标准键，避免前端因过滤异常看不到字段 + return { + "match_result": state.get("match_result", []), + "template_gremlin": state.get("result", ""), + "raw_gremlin": state.get("raw_result", ""), + "template_execution_result": state.get("template_exec_res", ""), + "raw_execution_result": state.get("raw_exec_res", ""), + } diff --git a/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py b/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py index e78aa6d58..694ca014d 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py @@ -24,15 +24,16 @@ class GraphIndex: def __init__( - self, - graph_url: Optional[str] = huge_settings.graph_url, - graph_name: Optional[str] = huge_settings.graph_name, - graph_user: Optional[str] = huge_settings.graph_user, - graph_pwd: Optional[str] = huge_settings.graph_pwd, - graph_space: Optional[str] = huge_settings.graph_space, + self, + graph_url: Optional[str] = huge_settings.graph_url, + graph_name: Optional[str] = huge_settings.graph_name, + graph_user: Optional[str] = huge_settings.graph_user, + graph_pwd: Optional[str] = huge_settings.graph_pwd, + graph_space: Optional[str] = huge_settings.graph_space, ): - self.client = PyHugeClient(url=graph_url, graph=graph_name, user=graph_user, pwd=graph_pwd, - graphspace=graph_space) + self.client = PyHugeClient( + url=graph_url, graph=graph_name, user=graph_user, pwd=graph_pwd, graphspace=graph_space + ) def clear_graph(self): self.client.gremlin().exec("g.V().drop()") diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py index 641ac6d6e..f85483185 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py @@ -37,7 +37,9 @@ def __init__(self, embed_dim: int = 1024): self.properties = [] @staticmethod - def from_index_file(dir_path: str, filename_prefix: str = None, record_miss: bool = True) -> "VectorIndex": + def from_index_file( + dir_path: str, filename_prefix: str = None, record_miss: bool = True + ) -> "VectorIndex": """Load index from files, supporting model-specific filenames. This method loads a Faiss index and its corresponding properties from a directory. @@ -47,13 +49,18 @@ def from_index_file(dir_path: str, filename_prefix: str = None, record_miss: boo matches the number of properties. """ index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + property_name = ( + f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + ) index_file = os.path.join(dir_path, index_name) properties_file = os.path.join(dir_path, property_name) miss_files = [f for f in [index_file, properties_file] if not os.path.exists(f)] if miss_files: if record_miss: - log.warning("Missing vector files: %s. \nNeed create a new one for it.", ", ".join(miss_files)) + log.warning( + "Missing vector files: %s. \nNeed create a new one for it.", + ", ".join(miss_files), + ) return VectorIndex() try: @@ -61,7 +68,9 @@ def from_index_file(dir_path: str, filename_prefix: str = None, record_miss: boo with open(properties_file, "rb") as f: properties = pkl.load(f) except (RuntimeError, pkl.UnpicklingError, OSError) as e: - log.error("Failed to load index files for model '%s': %s", filename_prefix or "default", e) + log.error( + "Failed to load index files for model '%s': %s", filename_prefix or "default", e + ) raise RuntimeError( f"Could not load index files for model '{filename_prefix or 'default'}'. " f"Original error ({type(e).__name__}): {e}" @@ -85,7 +94,9 @@ def to_index_file(self, dir_path: str, filename_prefix: str = None): os.makedirs(dir_path) index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + property_name = ( + f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + ) index_file = os.path.join(dir_path, index_name) properties_file = os.path.join(dir_path, property_name) faiss.write_index(self.index, index_file) @@ -115,7 +126,9 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: self.properties = [p for i, p in enumerate(self.properties) if i not in indices] return remove_num - def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9) -> List[Any]: + def search( + self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9 + ) -> List[Any]: if self.index.ntotal == 0: return [] @@ -129,7 +142,9 @@ def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0 results.append(deepcopy(self.properties[i])) log.debug("[✓] Add valid distance %s to results.", dist) else: - log.debug("[x] Distance %s >= threshold %s, ignore this result.", dist, dis_threshold) + log.debug( + "[x] Distance %s >= threshold %s, ignore this result.", dist, dis_threshold + ) return results @staticmethod @@ -140,7 +155,9 @@ def clean(dir_path: str, filename_prefix: str = None): If model_name is None, it targets the default files. """ index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + property_name = ( + f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + ) index_file = os.path.join(dir_path, index_name) properties_file = os.path.join(dir_path, property_name) diff --git a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py index 47c70e1a4..c73242012 100644 --- a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py +++ b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py @@ -26,6 +26,7 @@ # TODO: we could use middleware(AOP) in the future (dig out the lifecycle of gradio & fastapi) class UseTimeMiddleware(BaseHTTPMiddleware): """Middleware to add process time to response headers""" + def __init__(self, app): super().__init__(app) @@ -33,7 +34,7 @@ async def dispatch(self, request: Request, call_next): # TODO: handle time record for async task pool in gradio start_time = time.perf_counter() response = await call_next(request) - process_time = (time.perf_counter() - start_time) * 1000 # ms + process_time = (time.perf_counter() - start_time) * 1000 # ms unit = "ms" if process_time > 1000: process_time /= 1000 @@ -46,6 +47,6 @@ async def dispatch(self, request: Request, call_next): request.method, request.query_params, request.client.host, - request.url + request.url, ) return response diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py index db9b2f105..698b92837 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py @@ -32,9 +32,9 @@ class SimilarityMode(str, Enum): def similarity( - embedding1: Union[List[float], np.ndarray], - embedding2: Union[List[float], np.ndarray], - mode: SimilarityMode = SimilarityMode.DEFAULT, + embedding1: Union[List[float], np.ndarray], + embedding2: Union[List[float], np.ndarray], + mode: SimilarityMode = SimilarityMode.DEFAULT, ) -> float: """Get embedding similarity.""" if isinstance(embedding1, list): @@ -57,28 +57,22 @@ class BaseEmbedding(ABC): # TODO: replace all the usage by get_texts_embeddings() & remove it in the future @deprecated("Use get_texts_embeddings() instead in the future.") @abstractmethod - def get_text_embedding( - self, - text: str - ) -> List[float]: + def get_text_embedding(self, text: str) -> List[float]: """Comment""" @abstractmethod - def get_texts_embeddings( - self, - texts: List[str] - ) -> List[List[float]]: + def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch. - + This method should efficiently process multiple texts at once by leveraging the embedding model's batching capabilities, which is typically more efficient than processing texts individually. - + Parameters ---------- texts : List[str] A list of text strings to be embedded. - + Returns ------- List[List[float]] @@ -87,12 +81,9 @@ def get_texts_embeddings( """ @abstractmethod - async def async_get_texts_embeddings( - self, - texts: List[str] - ) -> List[List[float]]: + async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch asynchronously. - + This method should efficiently process multiple texts at once by leveraging the embedding model's batching capabilities, which is typically more efficient than processing texts individually. @@ -101,7 +92,7 @@ async def async_get_texts_embeddings( ---------- texts : List[str] A list of text strings to be embedded. - + Returns ------- List[List[float]] @@ -111,9 +102,9 @@ async def async_get_texts_embeddings( @staticmethod def similarity( - embedding1: Union[List[float], np.ndarray], - embedding2: Union[List[float], np.ndarray], - mode: SimilarityMode = SimilarityMode.DEFAULT, + embedding1: Union[List[float], np.ndarray], + embedding2: Union[List[float], np.ndarray], + mode: SimilarityMode = SimilarityMode.DEFAULT, ) -> float: """Get embedding similarity.""" if isinstance(embedding1, list): diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index f4026ad7f..d0e15f000 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -23,12 +23,12 @@ class OpenAIEmbedding: def __init__( - self, - model_name: str = "text-embedding-3-small", - api_key: Optional[str] = None, - api_base: Optional[str] = None + self, + model_name: str = "text-embedding-3-small", + api_key: Optional[str] = None, + api_base: Optional[str] = None, ): - api_key = api_key or '' + api_key = api_key or "" self.client = OpenAI(api_key=api_key, base_url=api_base) self.aclient = AsyncOpenAI(api_key=api_key, base_url=api_base) self.model_name = model_name @@ -38,21 +38,18 @@ def get_text_embedding(self, text: str) -> List[float]: response = self.client.embeddings.create(input=text, model=self.model_name) return response.data[0].embedding - def get_texts_embeddings( - self, - texts: List[str] - ) -> List[List[float]]: + def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch. - + This method efficiently processes multiple texts at once by leveraging OpenAI's batching capabilities, which is more efficient than processing texts individually. - + Parameters ---------- texts : List[str] A list of text strings to be embedded. - + Returns ------- List[List[float]] @@ -64,7 +61,7 @@ def get_texts_embeddings( async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch asynchronously. - + This method should efficiently process multiple texts at once by leveraging the embedding model's batching capabilities, which is typically more efficient than processing texts individually. @@ -73,7 +70,7 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] ---------- texts : List[str] A list of text strings to be embedded. - + Returns ------- List[List[float]] diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/base.py b/hugegraph-llm/src/hugegraph_llm/models/llms/base.py index c6bfa44a8..69c082690 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/base.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/base.py @@ -24,48 +24,48 @@ class BaseLLM(ABC): @abstractmethod def generate( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, ) -> str: """Comment""" @abstractmethod async def agenerate( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, ) -> str: """Comment""" @abstractmethod def generate_streaming( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, - on_token_callback: Optional[Callable] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, + on_token_callback: Optional[Callable] = None, ) -> Generator[str, None, None]: """Comment""" @abstractmethod async def agenerate_streaming( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, - on_token_callback: Optional[Callable] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, + on_token_callback: Optional[Callable] = None, ) -> AsyncGenerator[str, None]: """Comment""" @abstractmethod def num_tokens_from_string( - self, - string: str, + self, + string: str, ) -> str: """Given a string returns the number of tokens the given string consists of""" @abstractmethod def max_allowed_token_length( - self, + self, ) -> int: """Returns the maximum number of tokens the LLM can handle""" 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 7e1eaab68..9121fca09 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py @@ -173,8 +173,4 @@ def get_text2gql_llm(self): if __name__ == "__main__": client = LLMs().get_chat_llm() print(client.generate(prompt="What is the capital of China?")) - print( - client.generate( - messages=[{"role": "user", "content": "What is the capital of China?"}] - ) - ) + print(client.generate(messages=[{"role": "user", "content": "What is the capital of China?"}])) diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/litellm.py b/hugegraph-llm/src/hugegraph_llm/models/llms/litellm.py index b9cc0f19f..6f3c8129c 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/litellm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/litellm.py @@ -51,7 +51,7 @@ def __init__( @retry( stop=stop_after_attempt(2), wait=wait_exponential(multiplier=1, min=2, max=5), - retry=retry_if_exception_type((RateLimitError, BudgetExceededError, APIError)) + retry=retry_if_exception_type((RateLimitError, BudgetExceededError, APIError)), ) def generate( self, @@ -80,12 +80,12 @@ def generate( @retry( stop=stop_after_attempt(2), wait=wait_exponential(multiplier=1, min=2, max=5), - retry=retry_if_exception_type((RateLimitError, BudgetExceededError, APIError)) + retry=retry_if_exception_type((RateLimitError, BudgetExceededError, APIError)), ) async def agenerate( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, ) -> str: """Generate a response to the query messages/prompt asynchronously.""" if messages is None: diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py index 5354ba306..6d08ce8cd 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py @@ -28,6 +28,7 @@ class OllamaClient(BaseLLM): """LLM wrapper should take in a prompt and return a string.""" + def __init__(self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwargs): self.model = model self.client = ollama.Client(host=f"http://{host}:{port}", **kwargs) @@ -49,9 +50,9 @@ def generate( messages=messages, ) usage = { - "prompt_tokens": response['prompt_eval_count'], - "completion_tokens": response['eval_count'], - "total_tokens": response['prompt_eval_count'] + response['eval_count'], + "prompt_tokens": response["prompt_eval_count"], + "completion_tokens": response["eval_count"], + "total_tokens": response["prompt_eval_count"] + response["eval_count"], } log.info("Token usage: %s", json.dumps(usage)) return response["message"]["content"] @@ -61,9 +62,9 @@ def generate( @retry(tries=3, delay=1) async def agenerate( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, ) -> str: """Comment""" if messages is None: @@ -75,9 +76,9 @@ async def agenerate( messages=messages, ) usage = { - "prompt_tokens": response['prompt_eval_count'], - "completion_tokens": response['eval_count'], - "total_tokens": response['prompt_eval_count'] + response['eval_count'], + "prompt_tokens": response["prompt_eval_count"], + "completion_tokens": response["eval_count"], + "total_tokens": response["prompt_eval_count"] + response["eval_count"], } log.info("Token usage: %s", json.dumps(usage)) return response["message"]["content"] @@ -96,11 +97,7 @@ def generate_streaming( assert prompt is not None, "Messages or prompt must be provided." messages = [{"role": "user", "content": prompt}] - for chunk in self.client.chat( - model=self.model, - messages=messages, - stream=True - ): + for chunk in self.client.chat(model=self.model, messages=messages, stream=True): if not chunk["message"]: log.debug("Received empty chunk['message'] in streaming chunk: %s", chunk) continue @@ -122,9 +119,7 @@ async def agenerate_streaming( try: async_generator = await self.async_client.chat( - model=self.model, - messages=messages, - stream=True + model=self.model, messages=messages, stream=True ) async for chunk in async_generator: token = chunk.get("message", {}).get("content", "") @@ -135,7 +130,6 @@ async def agenerate_streaming( print(f"Retrying LLM call {e}") raise e - def num_tokens_from_string( self, string: str, diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py index 88cea3976..e1088c890 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py @@ -42,7 +42,7 @@ def __init__( max_tokens: int = 8092, temperature: float = 0.01, ) -> None: - api_key = api_key or '' + api_key = api_key or "" self.client = OpenAI(api_key=api_key, base_url=api_base) self.aclient = AsyncOpenAI(api_key=api_key, base_url=api_base) self.model = model_name @@ -186,7 +186,7 @@ async def agenerate_streaming( temperature=self.temperature, max_tokens=self.max_tokens, messages=messages, - stream=True + stream=True, ) async for chunk in completions: if not chunk.choices: diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py index 1710acfc2..3bf481ce2 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py @@ -31,16 +31,21 @@ def __init__( self.base_url = base_url self.model = model - def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int] = None) -> List[str]: + def get_rerank_lists( + self, query: str, documents: List[str], top_n: Optional[int] = None + ) -> List[str]: if not top_n: top_n = len(documents) - assert top_n <= len(documents), "'top_n' should be less than or equal to the number of documents" + assert top_n <= len( + documents + ), "'top_n' should be less than or equal to the number of documents" if top_n == 0: return [] url = self.base_url from pyhugegraph.utils.constants import Constants + headers = { "accept": Constants.HEADER_CONTENT_TYPE, "content-type": Constants.HEADER_CONTENT_TYPE, diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py index aa9f0c061..6136d61b4 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py @@ -32,5 +32,7 @@ def get_reranker(self): model=llm_settings.reranker_model, ) if self.reranker_type == "siliconflow": - return SiliconReranker(api_key=llm_settings.reranker_api_key, model=llm_settings.reranker_model) + return SiliconReranker( + api_key=llm_settings.reranker_api_key, model=llm_settings.reranker_model + ) raise Exception("Reranker type is not supported!") diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py index d63b0ba3d..e4a9b550a 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py @@ -29,10 +29,14 @@ def __init__( self.api_key = api_key self.model = model - def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int] = None) -> List[str]: + def get_rerank_lists( + self, query: str, documents: List[str], top_n: Optional[int] = None + ) -> List[str]: if not top_n: top_n = len(documents) - assert top_n <= len(documents), "'top_n' should be less than or equal to the number of documents" + assert top_n <= len( + documents + ), "'top_n' should be less than or equal to the number of documents" if top_n == 0: return [] @@ -48,6 +52,7 @@ def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int "top_n": top_n, } from pyhugegraph.utils.constants import Constants + headers = { "accept": Constants.HEADER_CONTENT_TYPE, "content-type": Constants.HEADER_CONTENT_TYPE, diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/gremlin_execute.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/gremlin_execute.py new file mode 100644 index 000000000..98fdcdd1b --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/gremlin_execute.py @@ -0,0 +1,68 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from typing import Any, Dict + +from PyCGraph import CStatus + +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.utils.hugegraph_utils import run_gremlin_query + + +def _ensure_limit(query: str, default_limit: int = 100) -> str: + if not query: + return query + q_lower = query.lower() + if "limit(" in q_lower: + return query + if any(token in q_lower for token in ["g.v(", ".v(", "g.e(", ".e("]): + return f"{query}.limit({default_limit})" + return query + + +class GremlinExecuteNode(BaseNode): + def node_init(self): + return CStatus() + + def operator_schedule(self, data_json: Dict[str, Any]): + # Read requested outputs from wk_input + requested = getattr(self.wk_input, "requested_outputs", None) or [] + need_template = "template_execution_result" in requested + need_raw = "raw_execution_result" in requested + + tmpl_q = data_json.get("result", "") + raw_q = data_json.get("raw_result", "") + + if need_template: + try: + safe_q = _ensure_limit(tmpl_q) + data_json["template_exec_res"] = run_gremlin_query(query=safe_q) + except Exception as exc: # pylint: disable=broad-except + data_json["template_exec_res"] = f"{exc}" + else: + data_json["template_exec_res"] = "" + + if need_raw: + try: + safe_q = _ensure_limit(raw_q) + data_json["raw_exec_res"] = run_gremlin_query(query=safe_q) + except Exception as exc: # pylint: disable=broad-except + data_json["raw_exec_res"] = f"{exc}" + else: + data_json["raw_exec_res"] = "" + + return data_json 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 71c490b20..84719d9eb 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py @@ -62,7 +62,7 @@ def node_init(self): return CStatus() def operator_schedule(self, data_json): - print(f"check data json {data_json}") + log.debug("SchemaNode input state: %s", data_json) if self.schema.startswith("{"): try: return self.check_schema.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py new file mode 100644 index 000000000..eb033d869 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -0,0 +1,49 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from typing import Any, Dict + +from PyCGraph import CStatus + +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.gremlin_example_index_query import GremlinExampleIndexQuery +from hugegraph_llm.models.embeddings.init_embedding import Embeddings + + +class GremlinExampleIndexQueryNode(BaseNode): + operator: GremlinExampleIndexQuery + + def node_init(self): + # Build operator (index lazy-loading handled in operator) + embedding = Embeddings().get_embedding() + example_num = getattr(self.wk_input, "example_num", None) + if not isinstance(example_num, int): + example_num = 2 + # Clamp to [0, 10] + example_num = max(0, min(10, example_num)) + self.operator = GremlinExampleIndexQuery(embedding=embedding, num_examples=example_num) + return CStatus() + + def operator_schedule(self, data_json: Dict[str, Any]): + # Ensure query is present in context; degrade gracefully if empty + query = getattr(self.wk_input, "query", "") or "" + data_json["query"] = query + if not query: + data_json["match_result"] = [] + return data_json + # Operator.run writes match_result into context + return self.operator.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py index a28b41346..7df2e68e7 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py @@ -61,16 +61,12 @@ def node_init(self): # few_shot_schema: already parsed dict or raw JSON string few_shot_schema = {} - fss_src = ( - self.wk_input.few_shot_schema if self.wk_input.few_shot_schema else None - ) + fss_src = self.wk_input.few_shot_schema if self.wk_input.few_shot_schema else None if fss_src: try: few_shot_schema = json.loads(fss_src) except json.JSONDecodeError as e: - return CStatus( - -1, f"Few Shot Schema is not in a valid JSON format: {e}" - ) + return CStatus(-1, f"Few Shot Schema is not in a valid JSON format: {e}") _context_payload = { "raw_texts": raw_texts, diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py new file mode 100644 index 000000000..ffbafbaf4 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py @@ -0,0 +1,70 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import json +from typing import Any, Dict, Optional + +from PyCGraph import CStatus + +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize +from hugegraph_llm.models.llms.init_llm import LLMs +from hugegraph_llm.config import prompt as prompt_cfg + + +def _stable_schema_string(state_json: Dict[str, Any]) -> str: + if "simple_schema" in state_json and state_json["simple_schema"] is not None: + return json.dumps(state_json["simple_schema"], ensure_ascii=False, sort_keys=True) + if "schema" in state_json and state_json["schema"] is not None: + return json.dumps(state_json["schema"], ensure_ascii=False, sort_keys=True) + return "" + + +class Text2GremlinNode(BaseNode): + operator: GremlinGenerateSynthesize + + def node_init(self): + # Select LLM + llm = LLMs().get_text2gql_llm() + # Serialize schema deterministically + state_json = self.context.to_json() + schema_str = _stable_schema_string(state_json) + # Prompt fallback + gremlin_prompt: Optional[str] = getattr(self.wk_input, "gremlin_prompt", None) + if gremlin_prompt is None or not str(gremlin_prompt).strip(): + gremlin_prompt = prompt_cfg.gremlin_generate_prompt + # Keep vertices/properties empty for now + self.operator = GremlinGenerateSynthesize( + llm=llm, + schema=schema_str, + vertices=None, + gremlin_prompt=gremlin_prompt, + ) + return CStatus() + + def operator_schedule(self, data_json: Dict[str, Any]): + # Ensure query exists in context; return empty if not provided + query = getattr(self.wk_input, "query", "") or "" + data_json["query"] = query + if not query: + data_json["result"] = "" + data_json["raw_result"] = "" + return data_json + # increase call count for observability + prev = data_json.get("call_count", 0) or 0 + data_json["call_count"] = prev + 1 + return self.operator.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index c1c742032..fc729c11e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -59,12 +59,8 @@ def _validate_schema(self, schema: Dict[str, Any]) -> None: check_type(schema, dict, "Input data is not a dictionary.") if "vertexlabels" not in schema or "edgelabels" not in schema: log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") - check_type( - schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." - ) - check_type( - schema["edgelabels"], list, "'edgelabels' in input data is not a list." - ) + check_type(schema["vertexlabels"], list, "'vertexlabels' in input data is not a list.") + check_type(schema["edgelabels"], list, "'edgelabels' in input data is not a list.") def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): property_labels = schema.get("propertykeys", []) @@ -82,19 +78,13 @@ def _process_vertex_labels( for vertex_label in schema["vertexlabels"]: self._validate_vertex_label(vertex_label) properties = vertex_label["properties"] - primary_keys = self._process_keys( - vertex_label, "primary_keys", properties[:1] - ) + primary_keys = self._process_keys(vertex_label, "primary_keys", properties[:1]) if len(primary_keys) == 0: log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") vertex_label["primary_keys"] = primary_keys - nullable_keys = self._process_keys( - vertex_label, "nullable_keys", properties[1:] - ) + nullable_keys = self._process_keys(vertex_label, "nullable_keys", properties[1:]) vertex_label["nullable_keys"] = nullable_keys - self._add_missing_properties( - properties, property_labels, property_label_set - ) + self._add_missing_properties(properties, property_labels, property_label_set) def _process_edge_labels( self, schema: Dict[str, Any], property_labels: list, property_label_set: set @@ -102,17 +92,13 @@ def _process_edge_labels( for edge_label in schema["edgelabels"]: self._validate_edge_label(edge_label) properties = edge_label.get("properties", []) - self._add_missing_properties( - properties, property_labels, property_label_set - ) + self._add_missing_properties(properties, property_labels, property_label_set) def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") if "name" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'name'.") - check_type( - vertex_label["name"], str, "'name' in vertex_label is not of correct type." - ) + check_type(vertex_label["name"], str, "'name' in vertex_label is not of correct type.") if "properties" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'properties'.") check_type( @@ -133,9 +119,7 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: log_and_raise( "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." ) - check_type( - edge_label["name"], str, "'name' in edge_label is not of correct type." - ) + check_type(edge_label["name"], str, "'name' in edge_label is not of correct type.") check_type( edge_label["source_label"], str, @@ -147,13 +131,9 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: "'target_label' in edge_label is not of correct type.", ) - def _process_keys( - self, label: Dict[str, Any], key_type: str, default_keys: list - ) -> list: + def _process_keys(self, label: Dict[str, Any], key_type: str, default_keys: list) -> list: keys = label.get(key_type, default_keys) - check_type( - keys, list, f"'{key_type}' in {label['name']} is not of correct type." - ) + check_type(keys, list, f"'{key_type}' in {label['name']} is not of correct type.") new_keys = [key for key in keys if key in label["properties"]] return new_keys diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py index 910de20d5..dc5b15e00 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py @@ -126,15 +126,20 @@ def _rerank_with_vertex_degree( reranker = Rerankers().get_reranker() try: vertex_rerank_res = [ - reranker.get_rerank_lists(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list + reranker.get_rerank_lists(query, vertex_degree) + [""] + for vertex_degree in vertex_degree_list ] except requests.exceptions.RequestException as e: - log.warning("Online reranker fails, automatically switches to local bleu method: %s", e) + log.warning( + "Online reranker fails, automatically switches to local bleu method: %s", e + ) self.method = "bleu" self.switch_to_bleu = True if self.method == "bleu": - vertex_rerank_res = [_bleu_rerank(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list] + vertex_rerank_res = [ + _bleu_rerank(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list + ] depth = len(vertex_degree_list) for result in results: @@ -144,7 +149,9 @@ def _rerank_with_vertex_degree( knowledge_with_degree[result] += [""] * (depth - len(knowledge_with_degree[result])) def sort_key(res: str) -> Tuple[int, ...]: - return tuple(vertex_rerank_res[i].index(knowledge_with_degree[res][i]) for i in range(depth)) + return tuple( + vertex_rerank_res[i].index(knowledge_with_degree[res][i]) for i in range(depth) + ) sorted_results = sorted(results, key=sort_key) return sorted_results[:topn] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py index 895a3795a..0d9967020 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py @@ -58,7 +58,8 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context["keywords"] = keywords from hugegraph_llm.utils.log import log - log.info("KEYWORDS: %s", context['keywords']) + + log.info("KEYWORDS: %s", context["keywords"]) return context def _filter_keywords( diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index 65c95db5e..330890b5d 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -191,7 +191,7 @@ def merge_dedup_rerank( method=rerank_method, near_neighbor_first=near_neighbor_first, custom_related_information=custom_related_information, - topk_return_results=topk_return_results + topk_return_results=topk_return_results, ) ) return self @@ -245,7 +245,7 @@ def run(self, **kwargs) -> Dict[str, Any]: """ if len(self._operators) == 0: self.extract_keywords().query_graphdb( - max_graph_items=kwargs.get('max_graph_items') + max_graph_items=kwargs.get("max_graph_items") ).synthesize_answer() context = kwargs 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 9eec04f7f..52626b72b 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 @@ -50,9 +50,7 @@ def run(self, data: dict) -> Dict[str, Any]: if not schema: # TODO: ensure the function works correctly (update the logic later) self.schema_free_mode(data.get("triples", [])) - log.warning( - "Using schema_free mode, could try schema_define mode for better effect!" - ) + log.warning("Using schema_free mode, could try schema_define mode for better effect!") else: self.init_schema_if_need(schema) self.load_into_graph(vertices, edges, schema) @@ -68,9 +66,7 @@ def _set_default_property(self, key, input_properties, property_label_map): # list or set default_value = [] input_properties[key] = default_value - log.warning( - "Property '%s' missing in vertex, set to '%s' for now", key, default_value - ) + log.warning("Property '%s' missing in vertex, set to '%s' for now", key, default_value) def _handle_graph_creation(self, func, *args, **kwargs): try: @@ -84,13 +80,9 @@ 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) - vertex_label_map = { - v_label["name"]: v_label for v_label in schema["vertexlabels"] - } + 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"] - } + property_label_map = {p_label["name"]: p_label for p_label in schema["propertykeys"]} for vertex in vertices: input_label = vertex["label"] @@ -106,9 +98,7 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- vertex_label = vertex_label_map[input_label] primary_keys = vertex_label["primary_keys"] nullable_keys = vertex_label.get("nullable_keys", []) - non_null_keys = [ - key for key in vertex_label["properties"] if key not in nullable_keys - ] + non_null_keys = [key for key in vertex_label["properties"] if key not in nullable_keys] has_problem = False # 2. Handle primary-keys mode vertex @@ -140,9 +130,7 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- # 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 - ) + 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(): @@ -179,9 +167,7 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- 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 - ) + self._handle_graph_creation(self.client.graph().addEdge, label, start, end, properties) def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -205,20 +191,18 @@ def init_schema_if_need(self, schema: dict): source_vertex_label = edge["source_label"] target_vertex_label = edge["target_label"] properties = edge["properties"] - self.schema.edgeLabel(edge_label).sourceLabel( - source_vertex_label - ).targetLabel(target_vertex_label).properties(*properties).nullableKeys( - *properties - ).ifNotExist().create() + self.schema.edgeLabel(edge_label).sourceLabel(source_vertex_label).targetLabel( + target_vertex_label + ).properties(*properties).nullableKeys(*properties).ifNotExist().create() def schema_free_mode(self, 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("name").ifNotExist().create() + self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel("vertex").properties( + "name" + ).ifNotExist().create() self.schema.indexLabel("vertexByName").onV("vertex").by( "name" @@ -278,9 +262,7 @@ def _set_property_data_type(self, property_key, data_type): log.warning("UUID type is not supported, use text instead") property_key.asText() else: - log.error( - "Unknown data type %s for property_key %s", data_type, property_key - ) + log.error("Unknown data type %s for property_key %s", data_type, property_key) def _set_property_cardinality(self, property_key, cardinality): if cardinality == PropertyCardinality.SINGLE: @@ -290,13 +272,9 @@ def _set_property_cardinality(self, property_key, cardinality): elif cardinality == PropertyCardinality.SET: property_key.valueSet() else: - log.error( - "Unknown cardinality %s for property_key %s", cardinality, property_key - ) + 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: + def _check_property_data_type(self, data_type: str, cardinality: str, value) -> bool: if cardinality in ( PropertyCardinality.LIST.value, PropertyCardinality.SET.value, @@ -326,9 +304,7 @@ def _check_single_data_type(self, data_type: str, value) -> bool: if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below - if ( - data_type == PropertyDataType.DATE.value - ): # the format should be "yyyy-MM-dd" + if data_type == PropertyDataType.DATE.value: # the format should be "yyyy-MM-dd" import re return isinstance(value, str) and re.match(r"^\d{4}-\d{2}-\d{2}$", value) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py index 6012b7534..bcff5f07b 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py @@ -196,8 +196,8 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: log.debug("Kneighbor gremlin query: %s", gremlin_query) paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = self._format_graph_query_result( - query_paths=paths + graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( + self._format_graph_query_result(query_paths=paths) ) # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) @@ -220,17 +220,21 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: max_deep=self._max_deep, max_items=self._max_items, ) - log.warning("Unable to find vid, downgraded to property query, please confirm if it meets expectation.") + log.warning( + "Unable to find vid, downgraded to property query, please confirm if it meets expectation." + ) paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = self._format_graph_query_result( - query_paths=paths + graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( + self._format_graph_query_result(query_paths=paths) ) context["graph_result"] = list(graph_chain_knowledge) if context["graph_result"]: context["graph_result_flag"] = 0 - context["vertex_degree_list"] = [list(vertex_degree) for vertex_degree in vertex_degree_list] + context["vertex_degree_list"] = [ + list(vertex_degree) for vertex_degree in vertex_degree_list + ] context["knowledge_with_degree"] = knowledge_with_degree context["graph_context_head"] = ( f"The following are graph knowledge in {self._max_deep} depth, e.g:\n" @@ -272,7 +276,9 @@ def _format_graph_from_vertex(self, query_result: List[Any]) -> Set[str]: knowledge.add(node_str) return knowledge - def _format_graph_query_result(self, query_paths) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: + def _format_graph_query_result( + self, query_paths + ) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: use_id_to_match = self._prop_to_match is None subgraph = set() subgraph_with_degree = {} @@ -282,7 +288,9 @@ def _format_graph_query_result(self, query_paths) -> Tuple[Set[str], List[Set[st for path in query_paths: # 1. Process each path - path_str, vertex_with_degree = self._process_path(path, use_id_to_match, v_cache, e_cache) + path_str, vertex_with_degree = self._process_path( + path, use_id_to_match, v_cache, e_cache + ) subgraph.add(path_str) subgraph_with_degree[path_str] = vertex_with_degree # 2. Update vertex degree list @@ -291,7 +299,11 @@ def _format_graph_query_result(self, query_paths) -> Tuple[Set[str], List[Set[st return subgraph, vertex_degree_list, subgraph_with_degree def _process_path( - self, path: Any, use_id_to_match: bool, v_cache: Set[str], e_cache: Set[Tuple[str, str, str]] + self, + path: Any, + use_id_to_match: bool, + v_cache: Set[str], + e_cache: Set[Tuple[str, str, str]], ) -> Tuple[str, List[str]]: flat_rel = "" raw_flat_rel = path["objects"] @@ -306,7 +318,14 @@ def _process_path( if i % 2 == 0: # Process each vertex flat_rel, prior_edge_str_len, depth = self._process_vertex( - item, flat_rel, node_cache, prior_edge_str_len, depth, nodes_with_degree, use_id_to_match, v_cache + item, + flat_rel, + node_cache, + prior_edge_str_len, + depth, + nodes_with_degree, + use_id_to_match, + v_cache, ) else: # Process each edge @@ -333,7 +352,9 @@ def _process_vertex( return flat_rel, prior_edge_str_len, depth node_cache.add(matched_str) - props_str = ", ".join(f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v) + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v + ) # TODO: we may remove label id or replace with label name if matched_str in v_cache: @@ -356,10 +377,14 @@ def _process_edge( use_id_to_match: bool, e_cache: Set[Tuple[str, str, str]], ) -> Tuple[str, int]: - props_str = ", ".join(f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v) + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v + ) props_str = f"{{{props_str}}}" if props_str else "" prev_matched_str = ( - raw_flat_rel[i - 1]["id"] if use_id_to_match else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] + raw_flat_rel[i - 1]["id"] + if use_id_to_match + else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] ) edge_key = (item["inV"], item["label"], item["outV"]) @@ -369,12 +394,16 @@ def _process_edge( else: edge_label = item["label"] - edge_str = f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" + edge_str = ( + f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" + ) path_str += edge_str prior_edge_str_len = len(edge_str) return path_str, prior_edge_str_len - def _update_vertex_degree_list(self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str]) -> None: + def _update_vertex_degree_list( + self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str] + ) -> None: for depth, node_str in enumerate(nodes_with_degree): if depth >= len(vertex_degree_list): vertex_degree_list.append(set()) @@ -384,8 +413,8 @@ def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: schema = self._get_graph_schema() vertex_props_str, edge_props_str = schema.split("\n")[:2] # TODO: rename to vertex (also need update in the schema) - vertex_props_str = vertex_props_str[len("Vertex properties: "):].strip("[").strip("]") - edge_props_str = edge_props_str[len("Edge properties: "):].strip("[").strip("]") + vertex_props_str = vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") + edge_props_str = edge_props_str[len("Edge properties: ") :].strip("[").strip("]") vertex_labels = self._extract_label_names(vertex_props_str) edge_labels = self._extract_label_names(edge_props_str) return vertex_labels, edge_labels 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 c4e2124c3..90f1c00ea 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 @@ -40,9 +40,7 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: new_vertex = { - key: vertex[key] - for key in ["id", "name", "properties"] - if key in vertex + key: vertex[key] for key in ["id", "name", "properties"] if key in vertex } mini_schema["vertexlabels"].append(new_vertex) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index 657baf68e..6d9f96214 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -23,17 +23,25 @@ from hugegraph_llm.config import resource_path, llm_settings, huge_settings from hugegraph_llm.indices.vector_index import VectorIndex from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel, get_filename_prefix, get_index_folder_name +from hugegraph_llm.utils.embedding_utils import ( + get_embeddings_parallel, + get_filename_prefix, + get_index_folder_name, +) # FIXME: we need keep the logic same with build_semantic_index.py class BuildGremlinExampleIndex: def __init__(self, embedding: BaseEmbedding, examples: List[Dict[str, str]]): - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "gremlin_examples")) self.examples = examples self.embedding = embedding - self.filename_prefix = get_filename_prefix(llm_settings.embedding_type, getattr(embedding, "model_name", None)) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(embedding, "model_name", None) + ) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: # !: We have assumed that self.example is not empty 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 4b7c4e3d4..5689a59ac 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 @@ -24,15 +24,23 @@ from hugegraph_llm.indices.vector_index import VectorIndex from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager -from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel, get_filename_prefix, get_index_folder_name +from hugegraph_llm.utils.embedding_utils import ( + get_embeddings_parallel, + get_filename_prefix, + get_index_folder_name, +) from hugegraph_llm.utils.log import log class BuildSemanticIndex: def __init__(self, embedding: BaseEmbedding): - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) - self.filename_prefix = get_filename_prefix(llm_settings.embedding_type, getattr(embedding, "model_name", None)) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(embedding, "model_name", None) + ) self.vid_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) self.embedding = embedding self.sm = SchemaManager(huge_settings.graph_name) @@ -42,27 +50,19 @@ def _extract_names(self, vertices: list[str]) -> list[str]: def run(self, context: Dict[str, Any]) -> Dict[str, Any]: vertexlabels = self.sm.schema.getSchema()["vertexlabels"] - all_pk_flag = all( - data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels - ) + all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) past_vids = self.vid_index.properties # TODO: We should build vid vector index separately, especially when the vertices may be very large - present_vids = context[ - "vertices" - ] # Warning: data truncated by fetch_graph_data.py + present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py removed_vids = set(past_vids) - set(present_vids) removed_num = self.vid_index.remove(removed_vids) added_vids = list(set(present_vids) - set(past_vids)) if added_vids: - vids_to_process = ( - self._extract_names(added_vids) if all_pk_flag else added_vids - ) - added_embeddings = asyncio.run( - get_embeddings_parallel(self.embedding, vids_to_process) - ) + vids_to_process = self._extract_names(added_vids) if all_pk_flag else added_vids + added_embeddings = asyncio.run(get_embeddings_parallel(self.embedding, 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.to_index_file(self.index_dir, self.filename_prefix) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py index 5cdad0316..f5fb823c5 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py @@ -41,9 +41,7 @@ def __init__(self, embedding: BaseEmbedding): self.filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(self.embedding, "model_name", None) ) - self.vector_index = VectorIndex.from_index_file( - self.index_dir, self.filename_prefix - ) + self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: if "chunks" not in context: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index 96d1a3833..b680f2ca3 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -26,7 +26,11 @@ from hugegraph_llm.indices.vector_index import VectorIndex, INDEX_FILE_NAME, PROPERTIES_FILE_NAME from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel, get_filename_prefix, get_index_folder_name +from hugegraph_llm.utils.embedding_utils import ( + get_embeddings_parallel, + get_filename_prefix, + get_index_folder_name, +) from hugegraph_llm.utils.log import log @@ -34,16 +38,25 @@ class GremlinExampleIndexQuery: def __init__(self, embedding: BaseEmbedding = None, num_examples: int = 1): self.embedding = embedding or Embeddings().get_embedding() self.num_examples = num_examples - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "gremlin_examples")) - self.filename_prefix = get_filename_prefix(llm_settings.embedding_type, - getattr(self.embedding, "model_name", None)) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(self.embedding, "model_name", None) + ) self._ensure_index_exists() self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) def _ensure_index_exists(self): - index_name = f"{self.filename_prefix}_{INDEX_FILE_NAME}" if self.filename_prefix else INDEX_FILE_NAME - props_name = f"{self.filename_prefix}_{PROPERTIES_FILE_NAME}" if self.filename_prefix else PROPERTIES_FILE_NAME + index_name = ( + f"{self.filename_prefix}_{INDEX_FILE_NAME}" if self.filename_prefix else INDEX_FILE_NAME + ) + props_name = ( + f"{self.filename_prefix}_{PROPERTIES_FILE_NAME}" + if self.filename_prefix + else PROPERTIES_FILE_NAME + ) if not ( os.path.exists(os.path.join(self.index_dir, index_name)) and os.path.exists(os.path.join(self.index_dir, props_name)) @@ -61,7 +74,9 @@ def _get_match_result(self, context: Dict[str, Any], query: str) -> List[Dict[st return self.vector_index.search(query_embedding, self.num_examples, dis_threshold=1.8) def _build_default_example_index(self): - properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict(orient="records") + properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict( + orient="records" + ) # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) queries = [row["query"] for row in properties] embeddings = asyncio.run(get_embeddings_parallel(self.embedding, queries)) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py index 8e195453d..3ac03246f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py @@ -31,16 +31,20 @@ class SemanticIdQuery: ID_QUERY_TEMPL = "g.V({vids_str}).limit(8)" def __init__( - self, - embedding: BaseEmbedding, - by: Literal["query", "keywords"] = "keywords", - topk_per_query: int = 10, - topk_per_keyword: int = huge_settings.topk_per_keyword, - vector_dis_threshold: float = huge_settings.vector_dis_threshold, + self, + embedding: BaseEmbedding, + by: Literal["query", "keywords"] = "keywords", + topk_per_query: int = 10, + topk_per_keyword: int = huge_settings.topk_per_keyword, + vector_dis_threshold: float = huge_settings.vector_dis_threshold, ): - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) - self.filename_prefix = get_filename_prefix(llm_settings.embedding_type, getattr(embedding, "model_name", None)) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(embedding, "model_name", None) + ) self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) self.embedding = embedding self.by = by @@ -65,7 +69,7 @@ def _exact_match_vids(self, keywords: List[str]) -> Tuple[List[str], List[str]]: vids_str = ",".join([f"'{vid}'" for vid in possible_vids]) resp = self._client.gremlin().exec(SemanticIdQuery.ID_QUERY_TEMPL.format(vids_str=vids_str)) - searched_vids = [v['id'] for v in resp['data']] + searched_vids = [v["id"] for v in resp["data"]] unsearched_keywords = set(keywords) for vid in searched_vids: @@ -79,10 +83,13 @@ def _fuzzy_match_vids(self, keywords: List[str]) -> List[str]: fuzzy_match_result = [] for keyword in keywords: keyword_vector = self.embedding.get_texts_embeddings([keyword])[0] - results = self.vector_index.search(keyword_vector, top_k=self.topk_per_keyword, - dis_threshold=float(self.vector_dis_threshold)) + results = self.vector_index.search( + keyword_vector, + top_k=self.topk_per_keyword, + dis_threshold=float(self.vector_dis_threshold), + ) if results: - fuzzy_match_result.extend(results[:self.topk_per_keyword]) + fuzzy_match_result.extend(results[: self.topk_per_keyword]) return fuzzy_match_result def run(self, context: Dict[str, Any]) -> Dict[str, Any]: @@ -92,7 +99,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query_vector = self.embedding.get_texts_embeddings([query])[0] results = self.vector_index.search(query_vector, top_k=self.topk_per_query) if results: - graph_query_list.update(results[:self.topk_per_query]) + graph_query_list.update(results[: self.topk_per_query]) else: # by keywords keywords = context.get("keywords", []) if not keywords: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py index e29f50a76..4ed616929 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py @@ -30,9 +30,13 @@ class VectorIndexQuery: def __init__(self, embedding: BaseEmbedding, topk: int = 3): self.embedding = embedding self.topk = topk - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) - self.filename_prefix = get_filename_prefix(llm_settings.embedding_type, getattr(embedding, "model_name", None)) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(embedding, "model_name", None) + ) self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py index 4348477f6..3b5c63103 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py @@ -37,7 +37,12 @@ class KgBuilder: - def __init__(self, llm: BaseLLM, embedding: Optional[BaseEmbedding] = None, graph: Optional[PyHugeClient] = None): + def __init__( + self, + llm: BaseLLM, + embedding: Optional[BaseEmbedding] = None, + graph: Optional[PyHugeClient] = None, + ): self.operators = [] self.llm = llm self.embedding = embedding @@ -69,7 +74,9 @@ def chunk_split( return self def extract_info( - self, example_prompt: Optional[str] = None, extract_type: Literal["triples", "property_graph"] = "triples" + self, + example_prompt: Optional[str] = None, + extract_type: Literal["triples", "property_graph"] = "triples", ): if extract_type == "triples": self.operators.append(InfoExtract(self.llm, example_prompt)) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py index 5c4ab5fd3..9138f9e9b 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py @@ -62,17 +62,26 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = (f"{context_head_str}\n" - f"{self._context_body}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) response = self._llm.generate(prompt=final_prompt) return {"answer": response} graph_result_context, vector_result_context = self.handle_vector_graph(context) - context = asyncio.run(self.async_generate(context, context_head_str, context_tail_str, - vector_result_context, graph_result_context)) + context = asyncio.run( + self.async_generate( + context, + context_head_str, + context_tail_str, + vector_result_context, + graph_result_context, + ) + ) return context def init_llm(self, context): @@ -95,7 +104,9 @@ def handle_vector_graph(self, context): vector_result_context = "No (vector)phrase related to the query." graph_result = context.get("graph_result") if graph_result: - graph_context_head = context.get("graph_context_head", "Knowledge from graphdb for the query:\n") + graph_context_head = context.get( + "graph_context_head", "Knowledge from graphdb for the query:\n" + ) graph_result_context = graph_context_head + "\n".join( f"{i + 1}. {res}" for i, res in enumerate(graph_result) ) @@ -108,11 +119,13 @@ async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[st context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = (f"{context_head_str}\n" - f"{self._context_body}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) response = self._llm.generate(prompt=final_prompt) yield {"answer": response} return @@ -120,45 +133,60 @@ async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[st graph_result_context, vector_result_context = self.handle_vector_graph(context) async for context in self.async_streaming_generate( - context, - context_head_str, - context_tail_str, - vector_result_context, - graph_result_context + context, context_head_str, context_tail_str, vector_result_context, graph_result_context ): yield context - async def async_generate(self, context: Dict[str, Any], context_head_str: str, - context_tail_str: str, vector_result_context: str, - graph_result_context: str): + async def async_generate( + self, + context: Dict[str, Any], + context_head_str: str, + context_tail_str: str, + vector_result_context: str, + graph_result_context: str, + ): # async_tasks stores the async tasks for different answer types async_tasks = {} if self._raw_answer: final_prompt = self._question async_tasks["raw_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) if self._vector_only_answer: - context_str = (f"{context_head_str}\n" - f"{vector_result_context}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" + f"{vector_result_context}\n" + f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) - async_tasks["vector_only_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) + async_tasks["vector_only_task"] = asyncio.create_task( + self._llm.agenerate(prompt=final_prompt) + ) if self._graph_only_answer: - context_str = (f"{context_head_str}\n" - f"{graph_result_context}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" + f"{graph_result_context}\n" + f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) - async_tasks["graph_only_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) + async_tasks["graph_only_task"] = asyncio.create_task( + self._llm.agenerate(prompt=final_prompt) + ) if self._graph_vector_answer: context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = (f"{context_head_str}\n" - f"{context_body_str}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_tasks["graph_vector_task"] = asyncio.create_task( self._llm.agenerate(prompt=final_prompt) ) @@ -167,7 +195,7 @@ async def async_generate(self, context: Dict[str, Any], context_head_str: str, "raw_task": "raw_answer", "vector_only_task": "vector_only_answer", "graph_only_task": "graph_only_answer", - "graph_vector_task": "graph_vector_answer" + "graph_vector_task": "graph_vector_answer", } for task_key, context_key in async_tasks_mapping.items(): @@ -176,66 +204,95 @@ async def async_generate(self, context: Dict[str, Any], context_head_str: str, context[context_key] = response log.debug("Query Answer: %s", response) - ops = sum([self._raw_answer, self._vector_only_answer, self._graph_only_answer, self._graph_vector_answer]) - context['call_count'] = context.get('call_count', 0) + ops + ops = sum( + [ + self._raw_answer, + self._vector_only_answer, + self._graph_only_answer, + self._graph_vector_answer, + ] + ) + context["call_count"] = context.get("call_count", 0) + ops return context - async def async_streaming_generate(self, context: Dict[str, Any], context_head_str: str, - context_tail_str: str, vector_result_context: str, - graph_result_context: str) -> AsyncGenerator[Dict[str, Any], None]: + async def async_streaming_generate( + self, + context: Dict[str, Any], + context_head_str: str, + context_tail_str: str, + vector_result_context: str, + graph_result_context: str, + ) -> AsyncGenerator[Dict[str, Any], None]: # async_tasks stores the async tasks for different answer types async_generators = [] auto_id = 0 if self._raw_answer: final_prompt = self._question async_generators.append( - self.__llm_generate_with_meta_info(task_id=auto_id, target_key="raw_answer", prompt=final_prompt) + self.__llm_generate_with_meta_info( + task_id=auto_id, target_key="raw_answer", prompt=final_prompt + ) ) auto_id += 1 if self._vector_only_answer: - context_str = (f"{context_head_str}\n" - f"{vector_result_context}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" + f"{vector_result_context}\n" + f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_generators.append( self.__llm_generate_with_meta_info( - task_id=auto_id, - target_key="vector_only_answer", - prompt=final_prompt + task_id=auto_id, target_key="vector_only_answer", prompt=final_prompt ) ) auto_id += 1 if self._graph_only_answer: - context_str = (f"{context_head_str}\n" - f"{graph_result_context}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" + f"{graph_result_context}\n" + f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_generators.append( - self.__llm_generate_with_meta_info(task_id=auto_id, target_key="graph_only_answer", prompt=final_prompt) + self.__llm_generate_with_meta_info( + task_id=auto_id, target_key="graph_only_answer", prompt=final_prompt + ) ) auto_id += 1 if self._graph_vector_answer: context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = (f"{context_head_str}\n" - f"{context_body_str}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_generators.append( self.__llm_generate_with_meta_info( - task_id=auto_id, - target_key="graph_vector_answer", - prompt=final_prompt + task_id=auto_id, target_key="graph_vector_answer", prompt=final_prompt ) ) auto_id += 1 - ops = sum([self._raw_answer, self._vector_only_answer, self._graph_only_answer, self._graph_vector_answer]) - context['call_count'] = context.get('call_count', 0) + ops + ops = sum( + [ + self._raw_answer, + self._vector_only_answer, + self._graph_only_answer, + self._graph_vector_answer, + ] + ) + context["call_count"] = context.get("call_count", 0) + ops async_tasks = [asyncio.create_task(anext(gen)) for gen in async_generators] while True: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py index 817065aa0..2ac2eafff 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py @@ -53,7 +53,8 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: extract_triples_by_regex(llm_output, data) print( f"LLM {self.__class__.__name__} input:{prompt} \n" - f" output: {llm_output} \n data: {data}") + f" output: {llm_output} \n data: {data}" + ) data["call_count"] = data.get("call_count", 0) + 1 return data diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index 11f0f6022..650834300 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -54,7 +54,8 @@ def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional example_strings = [] for example in examples: example_strings.append( - f"- query: {example['query']}\n" f"- gremlin:\n```gremlin\n{example['gremlin']}\n```" + f"- query: {example['query']}\n" + f"- gremlin:\n```gremlin\n{example['gremlin']}\n```" ) return "\n\n".join(example_strings) @@ -89,11 +90,17 @@ async def async_generate(self, context: Dict[str, Any]): vertices=self._format_vertices(vertices=self.vertices), properties=self._format_properties(properties=None), ) - async_tasks["initialized_answer"] = asyncio.create_task(self.llm.agenerate(prompt=init_prompt)) + async_tasks["initialized_answer"] = asyncio.create_task( + self.llm.agenerate(prompt=init_prompt) + ) raw_response = await async_tasks["raw_answer"] initialized_response = await async_tasks["initialized_answer"] - log.debug("Text2Gremlin with tmpl prompt:\n %s,\n LLM Response: %s", init_prompt, initialized_response) + log.debug( + "Text2Gremlin with tmpl prompt:\n %s,\n LLM Response: %s", + init_prompt, + initialized_response, + ) context["result"] = self._extract_response(response=initialized_response) context["raw_result"] = self._extract_response(response=raw_response) @@ -123,7 +130,11 @@ def sync_generate(self, context: Dict[str, Any]): ) initialized_response = self.llm.generate(prompt=init_prompt) - log.debug("Text2Gremlin with tmpl prompt:\n %s,\n LLM Response: %s", init_prompt, initialized_response) + log.debug( + "Text2Gremlin with tmpl prompt:\n %s,\n LLM Response: %s", + init_prompt, + initialized_response, + ) context["result"] = self._extract_response(response=initialized_response) context["raw_result"] = self._extract_response(response=raw_response) 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 571ffde51..8897e0fea 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 @@ -198,12 +198,8 @@ def valid(self, element_id: str, max_length: int = 256) -> bool: return True def _filter_long_id(self, graph) -> Dict[str, List[Any]]: - graph["vertices"] = [ - vertex for vertex in graph["vertices"] if self.valid(vertex["id"]) - ] + graph["vertices"] = [vertex for vertex in graph["vertices"] if self.valid(vertex["id"])] graph["edges"] = [ - edge - for edge in graph["edges"] - if self.valid(edge["start"]) and self.valid(edge["end"]) + edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"]) ] return graph diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index b1e3c7db9..425f2a70b 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -31,12 +31,12 @@ class KeywordExtract: def __init__( - self, - text: Optional[str] = None, - llm: Optional[BaseLLM] = None, - max_keywords: int = 5, - extract_template: Optional[str] = None, - language: str = "english", + self, + text: Optional[str] = None, + llm: Optional[BaseLLM] = None, + max_keywords: int = 5, + extract_template: Optional[str] = None, + language: str = "english", ): self._llm = llm self._query = text @@ -76,17 +76,17 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: return context def _extract_keywords_from_response( - self, - response: str, - lowercase: bool = True, - start_token: str = "", + self, + response: str, + lowercase: bool = True, + start_token: str = "", ) -> Set[str]: keywords = [] # use re.escape(start_token) if start_token contains special chars like */&/^ etc. - matches = re.findall(rf'{start_token}[^\n]+\n?', response) + matches = re.findall(rf"{start_token}[^\n]+\n?", response) for match in matches: - match = match[len(start_token):].strip() + match = match[len(start_token) :].strip() keywords.extend( k.lower() if lowercase else k for k in re.split(r"[,，]+", match) @@ -98,5 +98,7 @@ def _extract_keywords_from_response( for token in keywords: sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: - results.update(w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language)) + results.update( + w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language) + ) return results diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py index 82326f000..058d1bce9 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py @@ -52,11 +52,11 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: few_shot_example = self._load_few_shot_example(example_name) meta_prompt = prompt_tpl.generate_extract_prompt_template.format( - few_shot_text=few_shot_example.get('text', ''), - few_shot_prompt=few_shot_example.get('prompt', ''), + few_shot_text=few_shot_example.get("text", ""), + few_shot_prompt=few_shot_example.get("prompt", ""), user_text=source_text, user_scenario=scenario, - language=prompt_tpl.llm_settings.language + language=prompt_tpl.llm_settings.language, ) log.debug("Meta-prompt sent to LLM: %s", meta_prompt) generated_prompt = self.llm.generate(prompt=meta_prompt) 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 79fb33b4f..565d79023 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 @@ -67,9 +67,9 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: item_type = item["type"] if item_type == "vertex": label = item["label"] - non_nullable_keys = set( - properties_map[item_type][label]["properties"] - ).difference(set(properties_map[item_type][label]["nullable_keys"])) + non_nullable_keys = set(properties_map[item_type][label]["properties"]).difference( + set(properties_map[item_type][label]["nullable_keys"]) + ) for key in non_nullable_keys: if key not in item["properties"]: item["properties"][key] = "NULL" @@ -82,9 +82,7 @@ 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) -> None: self.llm = llm self.example_prompt = example_prompt self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name @@ -142,9 +140,7 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: and "vertices" in property_graph and "edges" in property_graph ): - log.critical( - "Invalid property graph format; expecting 'vertices' and 'edges'." - ) + log.critical("Invalid property graph format; expecting 'vertices' and 'edges'.") return items # Create sets for valid vertex and edge labels based on the schema @@ -154,9 +150,7 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: def process_items(item_list, valid_labels, item_type): for item in item_list: if not isinstance(item, dict): - log.warning( - "Invalid property graph item type '%s'.", type(item) - ) + log.warning("Invalid property graph item type '%s'.", type(item)) continue if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): log.warning("Invalid item keys '%s'.", item.keys()) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py index 53587381a..928948413 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py @@ -34,7 +34,9 @@ def __init__( ): self.llm = llm or LLMs().get_chat_llm() # TODO: use a basic format for it - self.schema_prompt = schema_prompt or """ + self.schema_prompt = ( + schema_prompt + or """ You are a Graph Schema Generator for Apache HugeGraph. Based on the following three parts of content, output a Schema JSON that complies with HugeGraph specifications: @@ -53,6 +55,7 @@ def __init__( - Ensure the schema follows HugeGraph specifications - Do not include comments or extra fields. """ + ) def _format_raw_texts(self, raw_texts: List[str]) -> str: return "\n".join([f"- {text}" for text in raw_texts]) @@ -86,18 +89,15 @@ def build_prompt( self, raw_texts: List[str], query_examples: List[Dict[str, str]], - few_shot_schema: Dict[str, Any] + few_shot_schema: Dict[str, Any], ) -> str: return self.schema_prompt.format( raw_texts=self._format_raw_texts(raw_texts), query_examples=self._format_query_examples(query_examples), - few_shot_schema=self._format_few_shot_schema(few_shot_schema) + few_shot_schema=self._format_few_shot_schema(few_shot_schema), ) - def run( - self, - context: Dict[str, Any] - ) -> Dict[str, Any]: + def run(self, context: Dict[str, Any]) -> Dict[str, Any]: """Generate schema from context containing raw_texts, query_examples and few_shot_schema. Args: diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 6d3418c00..f941098b1 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -33,6 +33,11 @@ class WkFlowInput(GParam): source_text: str = None # Original text scenario: str = None # Scenario description example_name: str = None # Example name + # Fields for Text2Gremlin + query: str = None + example_num: int = None + gremlin_prompt: str = None + requested_outputs: Optional[List[str]] = None def reset(self, _: CStatus) -> None: self.texts = None @@ -49,6 +54,11 @@ def reset(self, _: CStatus) -> None: self.source_text = None self.scenario = None self.example_name = None + # Text2Gremlin related configuration + self.query = None + self.example_num = None + self.gremlin_prompt = None + self.requested_outputs = None class WkFlowState(GParam): @@ -66,6 +76,12 @@ class WkFlowState(GParam): keywords_embeddings = None generated_extract_prompt: Optional[str] = None + # Fields for Text2Gremlin results + match_result: Optional[List[dict]] = None + result: Optional[str] = None + raw_result: Optional[str] = None + template_exec_res: Optional[Any] = None + raw_exec_res: Optional[Any] = None def setup(self): self.schema = None @@ -74,7 +90,7 @@ def setup(self): self.edges = None self.vertices = None self.triples = None - self.call_count = None + self.call_count = 0 self.keywords = None self.vector_result = None @@ -82,6 +98,12 @@ def setup(self): self.keywords_embeddings = None self.generated_extract_prompt = None + # Text2Gremlin results reset + self.match_result = [] + self.result = "" + self.raw_result = "" + self.template_exec_res = "" + self.raw_exec_res = "" return CStatus() @@ -94,11 +116,7 @@ def to_json(self): dict: A dictionary containing non-None instance members and their serialized values. """ # Only export instance attributes (excluding methods and class attributes) whose values are not None - return { - k: v - for k, v in self.__dict__.items() - if not k.startswith("_") and v is not None - } + return {k: v for k, v in self.__dict__.items() if not k.startswith("_") and v is not None} # Implement a method that assigns keys from data_json as WkFlowState member variables def assign_from_json(self, data_json: dict): diff --git a/hugegraph-llm/src/hugegraph_llm/utils/anchor.py b/hugegraph-llm/src/hugegraph_llm/utils/anchor.py index d5f687a94..4542a7fd9 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/anchor.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/anchor.py @@ -15,16 +15,17 @@ from pathlib import Path + def get_project_root() -> Path: """ Returns the Path object of the project root directory. - - The function searches for common project root indicators like pyproject.toml + + The function searches for common project root indicators like pyproject.toml or .git directory by traversing up the directory tree from the current file location. - + Returns: Path: The absolute path to the project root directory - + Raises: RuntimeError: If no project root indicators could be found """ diff --git a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py index b07de6f4b..2914c4b28 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py @@ -109,6 +109,7 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any: def with_task_id(func: Callable) -> Callable: def wrapper(*args: Any, **kwargs: Any) -> Any: import uuid + task_id = f"task_{str(uuid.uuid4())[:8]}" log.debug("New task created with id: %s", task_id) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py index 55e50eadd..b2f485cea 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py @@ -24,7 +24,9 @@ from hugegraph_llm.models.embeddings.base import BaseEmbedding -async def _get_batch_with_progress(embedding: BaseEmbedding, batch: list[str], pbar: tqdm) -> list[Any]: +async def _get_batch_with_progress( + embedding: BaseEmbedding, batch: list[str], pbar: tqdm +) -> list[Any]: result = await embedding.async_get_texts_embeddings(batch) pbar.update(1) return result @@ -58,10 +60,7 @@ async def get_embeddings_parallel(embedding: BaseEmbedding, vids: list[str]) -> embeddings = [] with tqdm(total=len(vid_batches)) as pbar: # Create tasks for each batch with progress bar updates - tasks = [ - _get_batch_with_progress(embedding, batch, pbar) - for batch in vid_batches - ] + tasks = [_get_batch_with_progress(embedding, batch, pbar) for batch in vid_batches] # Use asyncio.gather() to preserve order batch_results = await asyncio.gather(*tasks) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index ccace69f2..7b870033a 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -45,13 +45,9 @@ def get_graph_index_info(): def get_graph_index_info_old(): - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) graph_summary_info = builder.fetch_graph_data().run() - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(builder.embedding, "model_name", None) @@ -66,16 +62,12 @@ def get_graph_index_info_old(): def clean_all_graph_index(): - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None), ) - VectorIndex.clean( - str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix - ) + VectorIndex.clean(str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix) VectorIndex.clean( str(os.path.join(resource_path, folder_name, "gremlin_examples")), filename_prefix, @@ -107,18 +99,14 @@ def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: def extract_graph_origin(input_file, input_text, schema, example_prompt) -> str: texts = read_documents(input_file, input_text) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) if not schema: return "ERROR: please input with correct schema/format." error_message = parse_schema(schema, builder) if error_message: return error_message - builder.chunk_split(texts, "document", "zh").extract_info( - example_prompt, "property_graph" - ) + builder.chunk_split(texts, "document", "zh").extract_info(example_prompt, "property_graph") try: context = builder.run() @@ -168,9 +156,7 @@ def update_vid_embedding(): def update_vid_embedding_old(): - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) builder.fetch_graph_data().build_vertex_id_semantic_index() log.debug("Operators: %s", builder.operators) try: @@ -199,9 +185,7 @@ def import_graph_data_old(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: data_json = json.loads(data.strip()) log.debug("Import graph data: %s", data) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) if schema: error_message = parse_schema(schema, builder) if error_message: @@ -222,9 +206,7 @@ def import_graph_data_old(data: str, schema: str) -> Union[str, Dict[str, Any]]: def build_schema(input_text, query_example, few_shot): scheduler = SchedulerSingleton.get_instance() try: - return scheduler.schedule_flow( - "build_schema", input_text, query_example, few_shot - ) + return scheduler.schedule_flow("build_schema", input_text, query_example, few_shot) except (TypeError, ValueError) as e: raise gr.Error(f"Schema generation failed: {e}") @@ -257,9 +239,7 @@ def build_schema_old(input_text, query_example, few_shot): except json.JSONDecodeError as e: raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) try: schema = builder.build_schema().run(context) except Exception as e: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index 1d02b45d3..147c0074c 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -53,7 +53,9 @@ def init_hg_test_data(): schema = client.schema() schema.propertyKey("name").asText().ifNotExist().create() schema.propertyKey("birthDate").asText().ifNotExist().create() - schema.vertexLabel("Person").properties("name", "birthDate").useCustomizeStringId().ifNotExist().create() + schema.vertexLabel("Person").properties( + "name", "birthDate" + ).useCustomizeStringId().ifNotExist().create() schema.vertexLabel("Movie").properties("name").useCustomizeStringId().ifNotExist().create() schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() @@ -110,13 +112,13 @@ def backup_data(): files = { "vertices.json": f"g.V().limit({MAX_VERTICES})" - f".aggregate('vertices').count().as('count').select('count','vertices')", + f".aggregate('vertices').count().as('count').select('count','vertices')", "edges.json": f"g.E().limit({MAX_EDGES}).aggregate('edges').count().as('count').select('count','edges')", - "schema.json": client.schema().getSchema(_format="groovy") + "schema.json": client.schema().getSchema(_format="groovy"), } vertexlabels = client.schema().getSchema()["vertexlabels"] - all_pk_flag = all(data.get('id_strategy') == 'PRIMARY_KEY' for data in vertexlabels) + all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) for filename, query in files.items(): write_backup_file(client, backup_subdir, filename, query, all_pk_flag) @@ -137,14 +139,22 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): json.dump(data, f, ensure_ascii=False) elif filename == "vertices.json": data_full = client.gremlin().exec(query)["data"][0]["vertices"] - data = [{key: value for key, value in vertex.items() if key != "id"} - for vertex in data_full] if all_pk_flag else data_full + data = ( + [ + {key: value for key, value in vertex.items() if key != "id"} + for vertex in data_full + ] + if all_pk_flag + else data_full + ) json.dump(data, f, ensure_ascii=False) elif filename == "schema.json": data_full = query if isinstance(data_full, dict) and "schema" in data_full: groovy_filename = filename.replace(".json", ".groovy") - with open(os.path.join(backup_subdir, groovy_filename), "w", encoding="utf-8") as groovy_file: + with open( + os.path.join(backup_subdir, groovy_filename), "w", encoding="utf-8" + ) as groovy_file: groovy_file.write(str(data_full["schema"])) else: data = data_full @@ -171,7 +181,7 @@ def manage_backup_retention(): raise Exception("Failed to manage backup retention") from e -#TODO: In the path demo/rag_demo/configs_block.py, +# TODO: In the path demo/rag_demo/configs_block.py, # there is a function test_api_connection that is similar to this function, # but it is not straightforward to reuse def check_graph_db_connection(url: str, name: str, user: str, pwd: str, graph_space: str) -> bool: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/log.py b/hugegraph-llm/src/hugegraph_llm/utils/log.py index 7076869fd..b64017454 100755 --- a/hugegraph-llm/src/hugegraph_llm/utils/log.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/log.py @@ -31,7 +31,7 @@ log_level=INFO, logger_name="root", propagate_logs=True, - stdout_logging=True + stdout_logging=True, ) # Initialize custom logger diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 138b0d359..301a6bdab 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -50,9 +50,7 @@ def read_documents(input_file, input_text): texts.append(text) elif full_path.endswith(".pdf"): # TODO: support PDF file - raise gr.Error( - "PDF will be supported later! Try to upload text/docx now" - ) + raise gr.Error("PDF will be supported later! Try to upload text/docx now") else: raise gr.Error("Please input txt or docx file.") else: @@ -62,9 +60,7 @@ def read_documents(input_file, input_text): # pylint: disable=C0301 def get_vector_index_info(): - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) ) @@ -91,15 +87,11 @@ def get_vector_index_info(): def clean_vector_index(): - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) ) - VectorIndex.clean( - str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix - ) + VectorIndex.clean(str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix) gr.Info("Clean vector index successfully!") diff --git a/hugegraph-llm/src/tests/config/test_config.py b/hugegraph-llm/src/tests/config/test_config.py index 6c803135f..7f480befa 100644 --- a/hugegraph-llm/src/tests/config/test_config.py +++ b/hugegraph-llm/src/tests/config/test_config.py @@ -23,5 +23,6 @@ class TestConfig(unittest.TestCase): def test_config(self): import nltk from hugegraph_llm.config import resource_path + nltk.data.path.append(resource_path) nltk.data.find("corpora/stopwords") diff --git a/hugegraph-llm/src/tests/models/embeddings/test_openai_embedding.py b/hugegraph-llm/src/tests/models/embeddings/test_openai_embedding.py index b9ded0f6c..f7afd15c6 100644 --- a/hugegraph-llm/src/tests/models/embeddings/test_openai_embedding.py +++ b/hugegraph-llm/src/tests/models/embeddings/test_openai_embedding.py @@ -22,6 +22,7 @@ class TestOpenAIEmbedding(unittest.TestCase): def test_embedding_dimension(self): from hugegraph_llm.models.embeddings.openai import OpenAIEmbedding + embedding = OpenAIEmbedding(api_key="") result = embedding.get_text_embedding("hello world!") print(result) diff --git a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py index caabe2a8e..7ad914468 100644 --- a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py +++ b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py @@ -28,7 +28,10 @@ def test_generate(self): def test_stream_generate(self): ollama_client = OllamaClient(model="llama3:8b-instruct-fp16") + def on_token_callback(chunk): print(chunk, end="", flush=True) - ollama_client.generate_streaming(prompt="What is the capital of France?", - on_token_callback=on_token_callback) + + ollama_client.generate_streaming( + prompt="What is the capital of France?", on_token_callback=on_token_callback + ) diff --git a/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py b/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py index d20a198f2..317d02879 100644 --- a/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py +++ b/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py @@ -26,12 +26,7 @@ def setUp(self): def test_schema_check_with_valid_input(self): data = { - "vertexlabels": [ - { - "name": "person", - "properties": ["name", "age", "occupation"] - } - ], + "vertexlabels": [{"name": "person", "properties": ["name", "age", "occupation"]}], "edgelabels": [ { "name": "knows", @@ -41,7 +36,7 @@ def test_schema_check_with_valid_input(self): ], } check_schema = CheckSchema(data) - self.assertEqual(check_schema.run(), {'schema': data}) + self.assertEqual(check_schema.run(), {"schema": data}) def test_schema_check_with_invalid_input(self): data = "invalid input" diff --git a/hugegraph-llm/src/tests/operators/common_op/test_nltk_helper.py b/hugegraph-llm/src/tests/operators/common_op/test_nltk_helper.py index 5ad73ed6f..b557cfc1b 100644 --- a/hugegraph-llm/src/tests/operators/common_op/test_nltk_helper.py +++ b/hugegraph-llm/src/tests/operators/common_op/test_nltk_helper.py @@ -22,6 +22,7 @@ class TestNLTKHelper(unittest.TestCase): def test_stopwords(self): from hugegraph_llm.operators.common_op.nltk_helper import NLTKHelper + nltk_helper = NLTKHelper() stopwords = nltk_helper.stopwords() print(stopwords) diff --git a/hugegraph-python-client/src/pyhugegraph/api/auth.py b/hugegraph-python-client/src/pyhugegraph/api/auth.py index 90b3e98d0..ab7d66169 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/auth.py +++ b/hugegraph-python-client/src/pyhugegraph/api/auth.py @@ -84,9 +84,7 @@ def create_group(self, group_name, group_description=None) -> Optional[Dict]: return self._invoke_request(data=json.dumps(data)) @router.http("DELETE", "auth/groups/{group_id}") - def delete_group( - self, group_id # pylint: disable=unused-argument - ) -> Optional[Dict]: + def delete_group(self, group_id) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("PUT", "auth/groups/{group_id}") @@ -116,9 +114,7 @@ def grant_accesses(self, group_id, target_id, access_permission) -> Optional[Dic ) @router.http("DELETE", "auth/accesses/{access_id}") - def revoke_accesses( - self, access_id # pylint: disable=unused-argument - ) -> Optional[Dict]: + def revoke_accesses(self, access_id) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("PUT", "auth/accesses/{access_id}") @@ -130,9 +126,7 @@ def modify_accesses( return self._invoke_request(data=json.dumps(data)) @router.http("GET", "auth/accesses/{access_id}") - def get_accesses( - self, access_id # pylint: disable=unused-argument - ) -> Optional[Dict]: + def get_accesses(self, access_id) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "auth/accesses") @@ -205,9 +199,7 @@ def update_belong( return self._invoke_request(data=json.dumps(data)) @router.http("GET", "auth/belongs/{belong_id}") - def get_belong( - self, belong_id # pylint: disable=unused-argument - ) -> Optional[Dict]: + def get_belong(self, belong_id) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "auth/belongs") diff --git a/hugegraph-python-client/src/pyhugegraph/api/graph.py b/hugegraph-python-client/src/pyhugegraph/api/graph.py index 907e01a5b..4555eeda4 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/graph.py +++ b/hugegraph-python-client/src/pyhugegraph/api/graph.py @@ -141,9 +141,7 @@ def addEdges(self, input_data) -> Optional[List[EdgeData]]: def appendEdge( self, edge_id, properties # pylint: disable=unused-argument ) -> Optional[EdgeData]: - if response := self._invoke_request( - data=json.dumps({"properties": properties}) - ): + if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None @@ -151,16 +149,12 @@ def appendEdge( def eliminateEdge( self, edge_id, properties # pylint: disable=unused-argument ) -> Optional[EdgeData]: - if response := self._invoke_request( - data=json.dumps({"properties": properties}) - ): + if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None @router.http("GET", "graph/edges/{edge_id}") - def getEdgeById( - self, edge_id # pylint: disable=unused-argument - ) -> Optional[EdgeData]: + def getEdgeById(self, edge_id) -> Optional[EdgeData]: # pylint: disable=unused-argument if response := self._invoke_request(): return EdgeData(response) return None diff --git a/hugegraph-python-client/src/pyhugegraph/api/schema.py b/hugegraph-python-client/src/pyhugegraph/api/schema.py index 8b4f54cfe..7e8926678 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/schema.py +++ b/hugegraph-python-client/src/pyhugegraph/api/schema.py @@ -64,9 +64,7 @@ def indexLabel(self, name): return index_label @router.http("GET", "schema?format={_format}") - def getSchema( - self, _format: str = "json" # pylint: disable=unused-argument - ) -> Optional[Dict]: + def getSchema(self, _format: str = "json") -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "schema/propertykeys/{property_name}") @@ -84,9 +82,7 @@ def getPropertyKeys(self) -> Optional[List[PropertyKeyData]]: return None @router.http("GET", "schema/vertexlabels/{name}") - def getVertexLabel( - self, name # pylint: disable=unused-argument - ) -> Optional[VertexLabelData]: + def getVertexLabel(self, name) -> Optional[VertexLabelData]: # pylint: disable=unused-argument if response := self._invoke_request(): return VertexLabelData(response) log.error("VertexLabel not found: %s", str(response)) @@ -128,9 +124,7 @@ def getRelations(self) -> Optional[List[str]]: return None @router.http("GET", "schema/indexlabels/{name}") - def getIndexLabel( - self, name # pylint: disable=unused-argument - ) -> Optional[IndexLabelData]: + def getIndexLabel(self, name) -> Optional[IndexLabelData]: # pylint: disable=unused-argument if response := self._invoke_request(): return IndexLabelData(response) log.error("IndexLabel not found: %s", str(response)) diff --git a/hugegraph-python-client/src/pyhugegraph/api/schema_manage/index_label.py b/hugegraph-python-client/src/pyhugegraph/api/schema_manage/index_label.py index 252d487bd..acef8f968 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/schema_manage/index_label.py +++ b/hugegraph-python-client/src/pyhugegraph/api/schema_manage/index_label.py @@ -83,11 +83,13 @@ def ifNotExist(self) -> "IndexLabel": @decorator_create def create(self): dic = self._parameter_holder.get_dic() - data = {"name": dic["name"], - "base_type": dic["base_type"], - "base_value": dic["base_value"], - "index_type": dic["index_type"], - "fields": list(dic["fields"])} + data = { + "name": dic["name"], + "base_type": dic["base_type"], + "base_value": dic["base_value"], + "index_type": dic["index_type"], + "fields": list(dic["fields"]), + } path = "schema/indexlabels" self.clean_parameter_holder() if response := self._sess.request(path, "POST", data=json.dumps(data)): diff --git a/hugegraph-python-client/src/pyhugegraph/api/services.py b/hugegraph-python-client/src/pyhugegraph/api/services.py index e086ae13e..f353673db 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/services.py +++ b/hugegraph-python-client/src/pyhugegraph/api/services.py @@ -87,9 +87,7 @@ def list_services(self, graphspace: str): # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "/graphspaces/{graphspace}/services/{service}") - def get_service( - self, graphspace: str, service: str # pylint: disable=unused-argument - ): + def get_service(self, graphspace: str, service: str): # pylint: disable=unused-argument """ Retrieve the details of a specific service. @@ -112,9 +110,7 @@ def get_service( """ return self._invoke_request() - def delete_service( - self, graphspace: str, service: str # pylint: disable=unused-argument - ): + def delete_service(self, graphspace: str, service: str): # pylint: disable=unused-argument """ Delete a specific service within a graph space. diff --git a/hugegraph-python-client/src/pyhugegraph/api/traverser.py b/hugegraph-python-client/src/pyhugegraph/api/traverser.py index 628c3f4bd..72dddb07a 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/traverser.py +++ b/hugegraph-python-client/src/pyhugegraph/api/traverser.py @@ -26,33 +26,23 @@ class TraverserManager(HugeParamsBase): def k_out(self, source_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() - @router.http( - "GET", 'traversers/kneighbor?source="{source_id}"&max_depth={max_depth}' - ) + @router.http("GET", 'traversers/kneighbor?source="{source_id}"&max_depth={max_depth}') def k_neighbor(self, source_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() - @router.http( - "GET", 'traversers/sameneighbors?vertex="{vertex_id}"&other="{other_id}"' - ) + @router.http("GET", 'traversers/sameneighbors?vertex="{vertex_id}"&other="{other_id}"') def same_neighbors(self, vertex_id, other_id): # pylint: disable=unused-argument return self._invoke_request() - @router.http( - "GET", 'traversers/jaccardsimilarity?vertex="{vertex_id}"&other="{other_id}"' - ) - def jaccard_similarity( - self, vertex_id, other_id # pylint: disable=unused-argument - ): + @router.http("GET", 'traversers/jaccardsimilarity?vertex="{vertex_id}"&other="{other_id}"') + def jaccard_similarity(self, vertex_id, other_id): # pylint: disable=unused-argument return self._invoke_request() @router.http( "GET", 'traversers/shortestpath?source="{source_id}"&target="{target_id}"&max_depth={max_depth}', ) - def shortest_path( - self, source_id, target_id, max_depth # pylint: disable=unused-argument - ): + def shortest_path(self, source_id, target_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() @router.http( @@ -78,9 +68,7 @@ def weighted_shortest_path( "GET", 'traversers/singlesourceshortestpath?source="{source_id}"&max_depth={max_depth}', ) - def single_source_shortest_path( - self, source_id, max_depth # pylint: disable=unused-argument - ): + def single_source_shortest_path(self, source_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() @router.http("POST", "traversers/multinodeshortestpath") @@ -114,9 +102,17 @@ def multi_node_shortest_path( def paths(self, source_id, target_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() - @router.http("POST", 'traversers/paths') + @router.http("POST", "traversers/paths") def advanced_paths( - self, sources, targets, step, max_depth, nearest=True, capacity=10000000, limit=10, with_vertex=False + self, + sources, + targets, + step, + max_depth, + nearest=True, + capacity=10000000, + limit=10, + with_vertex=False, ): return self._invoke_request( data=json.dumps( @@ -133,7 +129,6 @@ def advanced_paths( ) ) - @router.http("POST", "traversers/customizedpaths") def customized_paths( self, sources, steps, sort_by="INCR", with_vertex=True, capacity=-1, limit=-1 @@ -152,9 +147,7 @@ def customized_paths( ) @router.http("POST", "traversers/templatepaths") - def template_paths( - self, sources, targets, steps, capacity=10000, limit=10, with_vertex=True - ): + def template_paths(self, sources, targets, steps, capacity=10000, limit=10, with_vertex=True): return self._invoke_request( data=json.dumps( { @@ -172,9 +165,7 @@ def template_paths( "GET", 'traversers/crosspoints?source="{source_id}"&target="{target_id}"&max_depth={max_depth}', ) - def crosspoints( - self, source_id, target_id, max_depth # pylint: disable=unused-argument - ): + def crosspoints(self, source_id, target_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() @router.http("POST", "traversers/customizedcrosspoints") diff --git a/hugegraph-python-client/src/pyhugegraph/client.py b/hugegraph-python-client/src/pyhugegraph/client.py index 3b0301321..c9f4d1027 100644 --- a/hugegraph-python-client/src/pyhugegraph/client.py +++ b/hugegraph-python-client/src/pyhugegraph/client.py @@ -53,7 +53,7 @@ def __init__( user: str, pwd: str, graphspace: Optional[str] = None, - timeout: Optional[tuple[float, float]] = None + timeout: Optional[tuple[float, float]] = None, ): self.cfg = HGraphConfig(url, user, pwd, graph, graphspace, timeout or (0.5, 15.0)) diff --git a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py index 4bb70dba5..d5cc0eb9d 100644 --- a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py +++ b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py @@ -26,15 +26,13 @@ schema = client.schema() schema.propertyKey("name").asText().ifNotExist().create() schema.propertyKey("birthDate").asText().ifNotExist().create() - schema.vertexLabel("Person").properties( - "name", "birthDate" - ).usePrimaryKeyId().primaryKeys("name").ifNotExist().create() - schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys( + schema.vertexLabel("Person").properties("name", "birthDate").usePrimaryKeyId().primaryKeys( "name" ).ifNotExist().create() - schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel( - "Movie" + schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys( + "name" ).ifNotExist().create() + schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() print(schema.getVertexLabels()) print(schema.getEdgeLabels()) @@ -47,9 +45,7 @@ p2 = g.addVertex("Person", {"name": "Robert De Niro", "birthDate": "1943-08-17"}) m1 = g.addVertex("Movie", {"name": "The Godfather"}) m2 = g.addVertex("Movie", {"name": "The Godfather Part II"}) - m3 = g.addVertex( - "Movie", {"name": "The Godfather Coda The Death of Michael Corleone"} - ) + m3 = g.addVertex("Movie", {"name": "The Godfather Coda The Death of Michael Corleone"}) # add Edge g.addEdge("ActedIn", p1.id, m1.id, {}) diff --git a/hugegraph-python-client/src/pyhugegraph/structure/property_key_data.py b/hugegraph-python-client/src/pyhugegraph/structure/property_key_data.py index 6fb7c36f0..ff50d9b2f 100644 --- a/hugegraph-python-client/src/pyhugegraph/structure/property_key_data.py +++ b/hugegraph-python-client/src/pyhugegraph/structure/property_key_data.py @@ -62,5 +62,7 @@ def userdata(self): return self.__user_data def __repr__(self): - res = f"name: {self.__name}, cardinality: {self.__cardinality}, data_type: {self.__data_type}" + res = ( + f"name: {self.__name}, cardinality: {self.__cardinality}, data_type: {self.__data_type}" + ) return res diff --git a/hugegraph-python-client/src/pyhugegraph/utils/huge_config.py b/hugegraph-python-client/src/pyhugegraph/utils/huge_config.py index 3f6d78b95..429c07c6b 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/huge_config.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/huge_config.py @@ -39,7 +39,7 @@ class HGraphConfig: def __post_init__(self): # Add URL prefix compatibility check - if self.url and not self.url.startswith('http'): + if self.url and not self.url.startswith("http"): self.url = f"http://{self.url}" if self.graphspace and self.graphspace.strip(): @@ -47,9 +47,7 @@ def __post_init__(self): else: try: - response = requests.get( - f"{self.url}/versions", timeout=0.5 - ) + response = requests.get(f"{self.url}/versions", timeout=0.5) core = response.json()["versions"]["core"] log.info( # pylint: disable=logging-fstring-interpolation f"Retrieved API version information from the server: {core}." @@ -71,4 +69,6 @@ def __post_init__(self): except Exception: # pylint: disable=broad-exception-caught exc_type, exc_value, tb = sys.exc_info() traceback.print_exception(exc_type, exc_value, tb) - log.warning("Failed to retrieve API version information from the server, reverting to default v1.") + log.warning( + "Failed to retrieve API version information from the server, reverting to default v1." + ) diff --git a/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py b/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py index 0db81ed1c..f4a38a418 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py @@ -81,9 +81,7 @@ def wrapper(self: "HGraphContext", *args: Any, **kwargs: Any) -> Any: route = RouterRegistry().routers.get(func.__qualname__) if route.request_func is None: - route.request_func = functools.partial( - self.session.request, method=method - ) + route.request_func = functools.partial(self.session.request, method=method) return func(self, *args, **kwargs) @@ -134,9 +132,7 @@ def wrapper(self: "HGraphContext", *args: Any, **kwargs: Any) -> Any: formatted_path = path # Use functools.partial to create a partial function for making requests - make_request = functools.partial( - self.session.request, formatted_path, method - ) + make_request = functools.partial(self.session.request, formatted_path, method) # Store the partial function on the instance setattr(self, f"_{func.__name__}_request", make_request) diff --git a/hugegraph-python-client/src/pyhugegraph/utils/log.py b/hugegraph-python-client/src/pyhugegraph/utils/log.py index b263d32d5..c6f6bd074 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/log.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/log.py @@ -138,7 +138,9 @@ def init_logger( def _cached_log_file(filename): """Cache the opened file object""" # Use 1K buffer if writing to cloud storage - with open(filename, "a", buffering=_determine_buffer_size(filename), encoding="utf-8") as file_io: + with open( + filename, "a", buffering=_determine_buffer_size(filename), encoding="utf-8" + ) as file_io: atexit.register(file_io.close) return file_io diff --git a/hugegraph-python-client/src/pyhugegraph/utils/util.py b/hugegraph-python-client/src/pyhugegraph/utils/util.py index 90f27c24a..56a135547 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/util.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/util.py @@ -44,7 +44,9 @@ def create_exception(response_content): def check_if_authorized(response): if response.status_code == 401: - raise NotAuthorizedError(f"Please check your username and password. {str(response.content)}") + raise NotAuthorizedError( + f"Please check your username and password. {str(response.content)}" + ) return True @@ -56,8 +58,12 @@ def check_if_success(response, error=None): req = response.request req_body = req.body if req.body else "Empty body" response_body = response.text if response.text else "Empty body" - log.error("Error-Client: Request URL: %s, Request Body: %s, Response Body: %s", - req.url, req_body, response_body) + log.error( + "Error-Client: Request URL: %s, Request Body: %s, Response Body: %s", + req.url, + req_body, + response_body, + ) raise error return True @@ -103,9 +109,14 @@ def __call__(self, response: requests.Response, method: str, path: str): details = "key 'exception' not found" req_body = response.request.body if response.request.body else "Empty body" - req_body = req_body.encode('utf-8').decode('unicode_escape') - log.error("%s: %s\n[Body]: %s\n[Server Exception]: %s", - method, str(e).encode('utf-8').decode('unicode_escape'), req_body, details) + req_body = req_body.encode("utf-8").decode("unicode_escape") + log.error( + "%s: %s\n[Body]: %s\n[Server Exception]: %s", + method, + str(e).encode("utf-8").decode("unicode_escape"), + req_body, + details, + ) if response.status_code == 404: raise NotFoundError(response.content) from e diff --git a/hugegraph-python-client/src/tests/api/test_auth.py b/hugegraph-python-client/src/tests/api/test_auth.py index d2d30cf26..10e6bad7f 100644 --- a/hugegraph-python-client/src/tests/api/test_auth.py +++ b/hugegraph-python-client/src/tests/api/test_auth.py @@ -98,9 +98,7 @@ def test_group_operations(self): self.assertEqual(group["group_name"], "test_group") # Modify the group - group = self.auth.modify_group( - group["id"], group_description="test_description" - ) + group = self.auth.modify_group(group["id"], group_description="test_description") self.assertEqual(group["group_description"], "test_description") # Delete the group @@ -135,9 +133,7 @@ def test_target_operations(self): [{"type": "VERTEX", "label": "person", "properties": {"city": "Shanghai"}}], ) # Verify the target was modified - self.assertEqual( - target["target_resources"][0]["properties"]["city"], "Shanghai" - ) + self.assertEqual(target["target_resources"][0]["properties"]["city"], "Shanghai") # Delete the target self.auth.delete_target(target["id"]) diff --git a/hugegraph-python-client/src/tests/api/test_version.py b/hugegraph-python-client/src/tests/api/test_version.py index 1d6325dfd..44c5f376c 100644 --- a/hugegraph-python-client/src/tests/api/test_version.py +++ b/hugegraph-python-client/src/tests/api/test_version.py @@ -42,7 +42,7 @@ def tearDown(self): def test_version(self): version = self.version.version() self.assertIsInstance(version, dict) - self.assertIn("version", version['versions']) - self.assertIn("core", version['versions']) - self.assertIn("gremlin", version['versions']) - self.assertIn("api", version['versions']) + self.assertIn("version", version["versions"]) + self.assertIn("core", version["versions"]) + self.assertIn("gremlin", version["versions"]) + self.assertIn("api", version["versions"]) diff --git a/hugegraph-python-client/src/tests/client_utils.py b/hugegraph-python-client/src/tests/client_utils.py index 63b6d0770..f711072b8 100644 --- a/hugegraph-python-client/src/tests/client_utils.py +++ b/hugegraph-python-client/src/tests/client_utils.py @@ -28,7 +28,11 @@ class ClientUtils: def __init__(self): self.client = PyHugeClient( - url=self.URL, user=self.USERNAME, pwd=self.PASSWORD, graph=self.GRAPH, graphspace=self.GRAPHSPACE + url=self.URL, + user=self.USERNAME, + pwd=self.PASSWORD, + graph=self.GRAPH, + graphspace=self.GRAPHSPACE, ) assert self.client is not None From 85e1296e74c9913a927d264dc7f863dcb390434a Mon Sep 17 00:00:00 2001 From: Linyu <94553312+weijinglin@users.noreply.github.com> Date: Tue, 30 Sep 2025 10:55:47 +0800 Subject: [PATCH 04/71] Refactor RAG Workflow: Modularize Flows, Add Streaming, and Improve Node Initialization (#51) --- .../hugegraph_llm/demo/rag_demo/rag_block.py | 223 ++++++++++-------- .../src/hugegraph_llm/flows/common.py | 26 ++ .../flows/rag_flow_graph_only.py | 153 ++++++++++++ .../flows/rag_flow_graph_vector.py | 158 +++++++++++++ .../src/hugegraph_llm/flows/rag_flow_raw.py | 99 ++++++++ .../flows/rag_flow_vector_only.py | 123 ++++++++++ .../src/hugegraph_llm/flows/scheduler.py | 63 +++++ .../src/hugegraph_llm/nodes/base_node.py | 3 + .../nodes/common_node/merge_rerank_node.py | 83 +++++++ .../nodes/document_node/chunk_split.py | 2 +- .../hugegraph_node/commit_to_hugegraph.py | 3 +- .../nodes/hugegraph_node/fetch_graph_data.py | 3 +- .../nodes/hugegraph_node/graph_query_node.py | 93 ++++++++ .../nodes/hugegraph_node/schema.py | 3 +- .../nodes/index_node/build_semantic_index.py | 3 +- .../nodes/index_node/build_vector_index.py | 3 +- .../index_node/gremlin_example_index_query.py | 13 +- .../index_node/semantic_id_query_node.py | 91 +++++++ .../nodes/index_node/vector_query_node.py | 74 ++++++ .../nodes/llm_node/answer_synthesize_node.py | 99 ++++++++ .../nodes/llm_node/extract_info.py | 2 +- .../nodes/llm_node/keyword_extract_node.py | 80 +++++++ .../nodes/llm_node/prompt_generate.py | 2 +- .../nodes/llm_node/schema_build.py | 2 +- .../nodes/llm_node/text2gremlin.py | 10 +- .../src/hugegraph_llm/state/ai_state.py | 106 ++++++++- .../hugegraph_llm/utils/graph_index_utils.py | 115 ++------- 27 files changed, 1411 insertions(+), 224 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index c93ec5739..5ff3df931 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -21,13 +21,12 @@ from typing import AsyncGenerator, Tuple, Literal, Optional import gradio as gr +from hugegraph_llm.flows.scheduler import SchedulerSingleton import pandas as pd from gradio.utils import NamedString -from hugegraph_llm.config import resource_path, prompt, huge_settings, llm_settings -from hugegraph_llm.operators.graph_rag_task import RAGPipeline +from hugegraph_llm.config import resource_path, prompt, llm_settings from hugegraph_llm.utils.decorators import with_task_id -from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize from hugegraph_llm.utils.log import log @@ -72,44 +71,51 @@ def rag_answer( gr.Warning("Please select at least one generate mode.") return "", "", "", "" - rag = RAGPipeline() - if vector_search: - rag.query_vector_index() - if graph_search: - rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid( - vector_dis_threshold=vector_dis_threshold, - topk_per_keyword=topk_per_keyword, - ).import_schema(huge_settings.graph_name).query_graphdb( - num_gremlin_generate_example=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - max_graph_items=max_graph_items, - ) - # TODO: add more user-defined search strategies - rag.merge_dedup_rerank( - graph_ratio=graph_ratio, - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - topk_return_results=topk_return_results, - ) - rag.synthesize_answer( - raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt - ) - + scheduler = SchedulerSingleton.get_instance() try: - context = rag.run( - verbose=True, + # Select workflow by mode to avoid fetching the wrong pipeline from the pool + if graph_vector_answer or (graph_only_answer and vector_only_answer): + flow_key = "rag_graph_vector" + elif vector_only_answer: + flow_key = "rag_vector_only" + elif graph_only_answer: + flow_key = "rag_graph_only" + elif raw_answer: + flow_key = "rag_raw" + else: + raise RuntimeError("Unsupported flow type") + + res = scheduler.schedule_flow( + flow_key, query=text, vector_search=vector_search, graph_search=graph_search, + raw_answer=raw_answer, + vector_only_answer=vector_only_answer, + graph_only_answer=graph_only_answer, + graph_vector_answer=graph_vector_answer, + graph_ratio=graph_ratio, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + answer_prompt=answer_prompt, + keywords_extract_prompt=keywords_extract_prompt, + gremlin_tmpl_num=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, max_graph_items=max_graph_items, + topk_return_results=topk_return_results, + vector_dis_threshold=vector_dis_threshold, + topk_per_keyword=topk_per_keyword, ) - if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") + if res.get("switch_to_bleu"): + gr.Warning( + "Online reranker fails, automatically switches to local bleu rerank." + ) return ( - context.get("raw_answer", ""), - context.get("vector_only_answer", ""), - context.get("graph_only_answer", ""), - context.get("graph_vector_answer", ""), + res.get("raw_answer", ""), + res.get("vector_only_answer", ""), + res.get("graph_only_answer", ""), + res.get("graph_vector_answer", ""), ) except ValueError as e: log.critical(e) @@ -187,44 +193,47 @@ async def rag_answer_streaming( yield "", "", "", "" return - rag = RAGPipeline() - if vector_search: - rag.query_vector_index() - if graph_search: - rag.extract_keywords( - extract_template=keywords_extract_prompt - ).keywords_to_vid().import_schema(huge_settings.graph_name).query_graphdb( - num_gremlin_generate_example=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - ) - rag.merge_dedup_rerank( - graph_ratio, - rerank_method, - near_neighbor_first, - ) - # rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) - try: - context = rag.run( - verbose=True, query=text, vector_search=vector_search, graph_search=graph_search - ) - if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") - answer_synthesize = AnswerSynthesize( + # Select the specific streaming workflow + scheduler = SchedulerSingleton.get_instance() + if graph_vector_answer or (graph_only_answer and vector_only_answer): + flow_key = "rag_graph_vector" + elif vector_only_answer: + flow_key = "rag_vector_only" + elif graph_only_answer: + flow_key = "rag_graph_only" + elif raw_answer: + flow_key = "rag_raw" + else: + raise RuntimeError("Unsupported flow type") + + async for res in scheduler.schedule_stream_flow( + flow_key, + query=text, + vector_search=vector_search, + graph_search=graph_search, raw_answer=raw_answer, vector_only_answer=vector_only_answer, graph_only_answer=graph_only_answer, graph_vector_answer=graph_vector_answer, - prompt_template=answer_prompt, - ) - async for context in answer_synthesize.run_streaming(context): - if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") + graph_ratio=graph_ratio, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + answer_prompt=answer_prompt, + keywords_extract_prompt=keywords_extract_prompt, + gremlin_tmpl_num=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, + ): + if res.get("switch_to_bleu"): + gr.Warning( + "Online reranker fails, automatically switches to local bleu rerank." + ) yield ( - context.get("raw_answer", ""), - context.get("vector_only_answer", ""), - context.get("graph_only_answer", ""), - context.get("graph_vector_answer", ""), + res.get("raw_answer", ""), + res.get("vector_only_answer", ""), + res.get("graph_only_answer", ""), + res.get("graph_vector_answer", ""), ) except ValueError as e: log.critical(e) @@ -242,7 +251,10 @@ def create_rag_block(): with gr.Column(scale=2): # with gr.Blocks().queue(max_size=20, default_concurrency_limit=5): inp = gr.Textbox( - value=prompt.default_question, label="Question", show_copy_button=True, lines=3 + value=prompt.default_question, + label="Question", + show_copy_button=True, + lines=3, ) # TODO: Only support inline formula now. Should support block formula @@ -272,7 +284,10 @@ def create_rag_block(): ) answer_prompt_input = gr.Textbox( - value=prompt.answer_prompt, label="Query Prompt", show_copy_button=True, lines=7 + value=prompt.answer_prompt, + label="Query Prompt", + show_copy_button=True, + lines=7, ) keywords_extract_prompt_input = gr.Textbox( value=prompt.keywords_extract_prompt, @@ -283,7 +298,9 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): - raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") + raw_radio = gr.Radio( + choices=[True, False], value=False, label="Basic LLM Answer" + ) vector_only_radio = gr.Radio( choices=[True, False], value=False, label="Vector-only Answer" ) @@ -307,7 +324,9 @@ def toggle_slider(enable): label="Rerank method", ) example_num = gr.Number( - value=-1, label="Template Num (<0 means disable text2gql) ", precision=0 + value=-1, + label="Template Num (<0 means disable text2gql) ", + precision=0, ) graph_ratio = gr.Slider( 0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False @@ -352,7 +371,7 @@ def toggle_slider(enable): """## 2. (Batch) Back-testing ) > 1. Download the template file & fill in the questions you want to test. > 2. Upload the file & click the button to generate answers. (Preview shows the first 40 lines) - > 3. The answer options are the same as the above RAG/Q&A frame + > 3. The answer options are the same as the above RAG/Q&A frame """ ) tests_df_headers = [ @@ -366,7 +385,9 @@ def toggle_slider(enable): # FIXME: "demo" might conflict with the graph name, it should be modified. answers_path = os.path.join(resource_path, "demo", "questions_answers.xlsx") questions_path = os.path.join(resource_path, "demo", "questions.xlsx") - questions_template_path = os.path.join(resource_path, "demo", "questions_template.xlsx") + questions_template_path = os.path.join( + resource_path, "demo", "questions_template.xlsx" + ) def read_file_to_excel(file: NamedString, line_count: Optional[int] = None): df = None @@ -413,20 +434,23 @@ def several_rag_answer( total_rows = len(df) for index, row in df.iterrows(): question = row.iloc[0] - basic_llm_answer, vector_only_answer, graph_only_answer, graph_vector_answer = ( - rag_answer( - question, - is_raw_answer, - is_vector_only_answer, - is_graph_only_answer, - is_graph_vector_answer, - graph_ratio_ui, - rerank_method_ui, - near_neighbor_first_ui, - custom_related_information_ui, - answer_prompt, - keywords_extract_prompt, - ) + ( + basic_llm_answer, + vector_only_answer, + graph_only_answer, + graph_vector_answer, + ) = rag_answer( + question, + is_raw_answer, + is_vector_only_answer, + is_graph_only_answer, + is_graph_vector_answer, + graph_ratio_ui, + rerank_method_ui, + near_neighbor_first_ui, + custom_related_information_ui, + answer_prompt, + keywords_extract_prompt, ) df.at[index, "Basic LLM Answer"] = basic_llm_answer df.at[index, "Vector-only Answer"] = vector_only_answer @@ -443,12 +467,18 @@ def several_rag_answer( file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" ) with gr.Column(): - test_template_file = os.path.join(resource_path, "demo", "questions_template.xlsx") + test_template_file = os.path.join( + resource_path, "demo", "questions_template.xlsx" + ) gr.File(value=test_template_file, label="Download Template File") - answer_max_line_count = gr.Number(1, label="Max Lines To Show", minimum=1, maximum=40) + answer_max_line_count = gr.Number( + 1, label="Max Lines To Show", minimum=1, maximum=40 + ) answers_btn = gr.Button("Generate Answer (Batch)", variant="primary") # TODO: Set individual progress bars for dataframe - qa_dataframe = gr.DataFrame(label="Questions & Answers (Preview)", headers=tests_df_headers) + qa_dataframe = gr.DataFrame( + label="Questions & Answers (Preview)", headers=tests_df_headers + ) answers_btn.click( several_rag_answer, inputs=[ @@ -466,6 +496,15 @@ def several_rag_answer( ], outputs=[qa_dataframe, gr.File(label="Download Answered File", min_width=40)], ) - questions_file.change(read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count]) - answer_max_line_count.change(change_showing_excel, answer_max_line_count, qa_dataframe) - return inp, answer_prompt_input, keywords_extract_prompt_input, custom_related_information + questions_file.change( + read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count] + ) + answer_max_line_count.change( + change_showing_excel, answer_max_line_count, qa_dataframe + ) + return ( + inp, + answer_prompt_input, + keywords_extract_prompt_input, + custom_related_information, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/common.py b/hugegraph-llm/src/hugegraph_llm/flows/common.py index 4c552626a..e2348466c 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/common.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/common.py @@ -14,8 +14,10 @@ # limitations under the License. from abc import ABC, abstractmethod +from typing import Dict, Any, AsyncGenerator from hugegraph_llm.state.ai_state import WkFlowInput +from hugegraph_llm.utils.log import log class BaseFlow(ABC): @@ -43,3 +45,27 @@ def post_deal(self, *args, **kwargs): Post-processing interface. """ pass + + async def post_deal_stream( + self, pipeline=None + ) -> AsyncGenerator[Dict[str, Any], None]: + """ + Streaming post-processing interface. + Subclasses can override this method as needed. + """ + flow_name = self.__class__.__name__ + if pipeline is None: + yield {"error": "No pipeline provided"} + return + try: + state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info(f"{flow_name} post processing success") + stream_flow = state_json.get("stream_generator") + if stream_flow is None: + yield {"error": "No stream_generator found in workflow state"} + return + async for chunk in stream_flow: + yield chunk + except Exception as e: + log.error(f"{flow_name} post processing failed: {e}") + yield {"error": f"Post processing failed: {str(e)}"} diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py new file mode 100644 index 000000000..5feb3d471 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py @@ -0,0 +1,153 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +from typing import Optional, Literal + +from PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.llm_node.keyword_extract_node import KeywordExtractNode +from hugegraph_llm.nodes.index_node.semantic_id_query_node import SemanticIdQueryNode +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.nodes.hugegraph_node.graph_query_node import GraphQueryNode +from hugegraph_llm.nodes.common_node.merge_rerank_node import MergeRerankNode +from hugegraph_llm.nodes.llm_node.answer_synthesize_node import AnswerSynthesizeNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.utils.log import log + + +class RAGGraphOnlyFlow(BaseFlow): + """ + Workflow for graph-only answering (graph_only_answer) + """ + + def prepare( + self, + prepared_input: WkFlowInput, + query: str, + vector_search: bool = None, + graph_search: bool = None, + raw_answer: bool = None, + vector_only_answer: bool = None, + graph_only_answer: bool = None, + graph_vector_answer: bool = None, + graph_ratio: float = 0.5, + rerank_method: Literal["bleu", "reranker"] = "bleu", + near_neighbor_first: bool = False, + custom_related_information: str = "", + answer_prompt: Optional[str] = None, + keywords_extract_prompt: Optional[str] = None, + gremlin_tmpl_num: Optional[int] = -1, + gremlin_prompt: Optional[str] = None, + max_graph_items: int = None, + topk_return_results: int = None, + vector_dis_threshold: float = None, + topk_per_keyword: int = None, + **_: dict, + ): + prepared_input.query = query + prepared_input.vector_search = vector_search + prepared_input.graph_search = graph_search + prepared_input.raw_answer = raw_answer + prepared_input.vector_only_answer = vector_only_answer + prepared_input.graph_only_answer = graph_only_answer + prepared_input.graph_vector_answer = graph_vector_answer + prepared_input.gremlin_tmpl_num = gremlin_tmpl_num + prepared_input.gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt + prepared_input.max_graph_items = ( + max_graph_items or huge_settings.max_graph_items + ) + prepared_input.topk_per_keyword = ( + topk_per_keyword or huge_settings.topk_per_keyword + ) + prepared_input.topk_return_results = ( + topk_return_results or huge_settings.topk_return_results + ) + prepared_input.rerank_method = rerank_method + prepared_input.near_neighbor_first = near_neighbor_first + prepared_input.keywords_extract_prompt = ( + keywords_extract_prompt or prompt.keywords_extract_prompt + ) + prepared_input.answer_prompt = answer_prompt or prompt.answer_prompt + prepared_input.custom_related_information = custom_related_information + prepared_input.vector_dis_threshold = ( + vector_dis_threshold or huge_settings.vector_dis_threshold + ) + prepared_input.schema = huge_settings.graph_name + + prepared_input.data_json = { + "query": query, + "vector_search": vector_search, + "graph_search": graph_search, + "max_graph_items": max_graph_items or huge_settings.max_graph_items, + } + return + + def build_flow(self, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, **kwargs) + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + # Create nodes and register them with registerGElement + only_keyword_extract_node = KeywordExtractNode() + only_semantic_id_query_node = SemanticIdQueryNode() + only_schema_node = SchemaNode() + only_graph_query_node = GraphQueryNode() + merge_rerank_node = MergeRerankNode() + answer_synthesize_node = AnswerSynthesizeNode() + + pipeline.registerGElement(only_keyword_extract_node, set(), "only_keyword") + pipeline.registerGElement( + only_semantic_id_query_node, {only_keyword_extract_node}, "only_semantic" + ) + pipeline.registerGElement(only_schema_node, set(), "only_schema") + pipeline.registerGElement( + only_graph_query_node, + {only_schema_node, only_semantic_id_query_node}, + "only_graph", + ) + pipeline.registerGElement( + merge_rerank_node, {only_graph_query_node}, "merge_one" + ) + pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "graph") + log.info("RAGGraphOnlyFlow pipeline built successfully") + return pipeline + + def post_deal(self, pipeline=None): + if pipeline is None: + return json.dumps( + {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 + ) + try: + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGGraphOnlyFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } + except Exception as e: + log.error(f"RAGGraphOnlyFlow post processing failed: {e}") + return json.dumps( + {"error": f"Post processing failed: {str(e)}"}, + ensure_ascii=False, + indent=2, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py new file mode 100644 index 000000000..2f4a2bfa2 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py @@ -0,0 +1,158 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +from typing import Optional, Literal + +from PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.index_node.vector_query_node import VectorQueryNode +from hugegraph_llm.nodes.llm_node.keyword_extract_node import KeywordExtractNode +from hugegraph_llm.nodes.index_node.semantic_id_query_node import SemanticIdQueryNode +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.nodes.hugegraph_node.graph_query_node import GraphQueryNode +from hugegraph_llm.nodes.common_node.merge_rerank_node import MergeRerankNode +from hugegraph_llm.nodes.llm_node.answer_synthesize_node import AnswerSynthesizeNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.utils.log import log + + +class RAGGraphVectorFlow(BaseFlow): + """ + Workflow for graph + vector hybrid answering (graph_vector_answer) + """ + + def prepare( + self, + prepared_input: WkFlowInput, + query: str, + vector_search: bool = None, + graph_search: bool = None, + raw_answer: bool = None, + vector_only_answer: bool = None, + graph_only_answer: bool = None, + graph_vector_answer: bool = None, + graph_ratio: float = 0.5, + rerank_method: Literal["bleu", "reranker"] = "bleu", + near_neighbor_first: bool = False, + custom_related_information: str = "", + answer_prompt: Optional[str] = None, + keywords_extract_prompt: Optional[str] = None, + gremlin_tmpl_num: Optional[int] = -1, + gremlin_prompt: Optional[str] = None, + max_graph_items: int = None, + topk_return_results: int = None, + vector_dis_threshold: float = None, + topk_per_keyword: int = None, + **_: dict, + ): + prepared_input.query = query + prepared_input.vector_search = vector_search + prepared_input.graph_search = graph_search + prepared_input.raw_answer = raw_answer + prepared_input.vector_only_answer = vector_only_answer + prepared_input.graph_only_answer = graph_only_answer + prepared_input.graph_vector_answer = graph_vector_answer + prepared_input.graph_ratio = graph_ratio + prepared_input.gremlin_tmpl_num = gremlin_tmpl_num + prepared_input.gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt + prepared_input.max_graph_items = ( + max_graph_items or huge_settings.max_graph_items + ) + prepared_input.topk_return_results = ( + topk_return_results or huge_settings.topk_return_results + ) + prepared_input.topk_per_keyword = ( + topk_per_keyword or huge_settings.topk_per_keyword + ) + prepared_input.vector_dis_threshold = ( + vector_dis_threshold or huge_settings.vector_dis_threshold + ) + prepared_input.rerank_method = rerank_method + prepared_input.near_neighbor_first = near_neighbor_first + prepared_input.keywords_extract_prompt = ( + keywords_extract_prompt or prompt.keywords_extract_prompt + ) + prepared_input.answer_prompt = answer_prompt or prompt.answer_prompt + prepared_input.custom_related_information = custom_related_information + prepared_input.schema = huge_settings.graph_name + + prepared_input.data_json = { + "query": query, + "vector_search": vector_search, + "graph_search": graph_search, + "max_graph_items": max_graph_items or huge_settings.max_graph_items, + } + return + + def build_flow(self, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, **kwargs) + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + # Create nodes (registration style consistent with RAGFlow) + vector_query_node = VectorQueryNode() + keyword_extract_node = KeywordExtractNode() + semantic_id_query_node = SemanticIdQueryNode() + schema_node = SchemaNode() + graph_query_node = GraphQueryNode() + merge_rerank_node = MergeRerankNode() + answer_synthesize_node = AnswerSynthesizeNode() + + # Register nodes and their dependencies + pipeline.registerGElement(vector_query_node, set(), "vector") + pipeline.registerGElement(keyword_extract_node, set(), "keyword") + pipeline.registerGElement( + semantic_id_query_node, {keyword_extract_node}, "semantic" + ) + pipeline.registerGElement(schema_node, set(), "schema") + pipeline.registerGElement( + graph_query_node, {schema_node, semantic_id_query_node}, "graph" + ) + pipeline.registerGElement( + merge_rerank_node, {graph_query_node, vector_query_node}, "merge" + ) + pipeline.registerGElement( + answer_synthesize_node, {merge_rerank_node}, "graph_vector" + ) + log.info("RAGGraphVectorFlow pipeline built successfully") + return pipeline + + def post_deal(self, pipeline=None): + if pipeline is None: + return json.dumps( + {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 + ) + try: + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGGraphVectorFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } + except Exception as e: + log.error(f"RAGGraphVectorFlow post processing failed: {e}") + return json.dumps( + {"error": f"Post processing failed: {str(e)}"}, + ensure_ascii=False, + indent=2, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py new file mode 100644 index 000000000..f62e574bb --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py @@ -0,0 +1,99 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +from typing import Optional + +from PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.llm_node.answer_synthesize_node import AnswerSynthesizeNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.utils.log import log + + +class RAGRawFlow(BaseFlow): + """ + Workflow for basic LLM answering only (raw_answer) + """ + + def prepare( + self, + prepared_input: WkFlowInput, + query: str, + vector_search: bool = None, + graph_search: bool = None, + raw_answer: bool = None, + vector_only_answer: bool = None, + graph_only_answer: bool = None, + graph_vector_answer: bool = None, + custom_related_information: str = "", + answer_prompt: Optional[str] = None, + max_graph_items: int = None, + **_: dict, + ): + prepared_input.query = query + prepared_input.raw_answer = raw_answer + prepared_input.vector_only_answer = vector_only_answer + prepared_input.graph_only_answer = graph_only_answer + prepared_input.graph_vector_answer = graph_vector_answer + prepared_input.custom_related_information = custom_related_information + prepared_input.answer_prompt = answer_prompt or prompt.answer_prompt + prepared_input.schema = huge_settings.graph_name + + prepared_input.data_json = { + "query": query, + "vector_search": vector_search, + "graph_search": graph_search, + "max_graph_items": max_graph_items or huge_settings.max_graph_items, + } + return + + def build_flow(self, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, **kwargs) + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + # Create nodes and register with registerGElement (no GRegion required) + answer_synthesize_node = AnswerSynthesizeNode() + pipeline.registerGElement(answer_synthesize_node, set(), "raw") + log.info("RAGRawFlow pipeline built successfully") + return pipeline + + def post_deal(self, pipeline=None): + if pipeline is None: + return json.dumps( + {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 + ) + try: + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGRawFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } + except Exception as e: + log.error(f"RAGRawFlow post processing failed: {e}") + return json.dumps( + {"error": f"Post processing failed: {str(e)}"}, + ensure_ascii=False, + indent=2, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py new file mode 100644 index 000000000..c727eacce --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py @@ -0,0 +1,123 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +from typing import Optional, Literal + +from PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.index_node.vector_query_node import VectorQueryNode +from hugegraph_llm.nodes.common_node.merge_rerank_node import MergeRerankNode +from hugegraph_llm.nodes.llm_node.answer_synthesize_node import AnswerSynthesizeNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.utils.log import log + + +class RAGVectorOnlyFlow(BaseFlow): + """ + Workflow for vector-only answering (vector_only_answer) + """ + + def prepare( + self, + prepared_input: WkFlowInput, + query: str, + vector_search: bool = None, + graph_search: bool = None, + raw_answer: bool = None, + vector_only_answer: bool = None, + graph_only_answer: bool = None, + graph_vector_answer: bool = None, + rerank_method: Literal["bleu", "reranker"] = "bleu", + near_neighbor_first: bool = False, + custom_related_information: str = "", + answer_prompt: Optional[str] = None, + max_graph_items: int = None, + topk_return_results: int = None, + vector_dis_threshold: float = None, + **_: dict, + ): + prepared_input.query = query + prepared_input.vector_search = vector_search + prepared_input.graph_search = graph_search + prepared_input.raw_answer = raw_answer + prepared_input.vector_only_answer = vector_only_answer + prepared_input.graph_only_answer = graph_only_answer + prepared_input.graph_vector_answer = graph_vector_answer + prepared_input.vector_dis_threshold = ( + vector_dis_threshold or huge_settings.vector_dis_threshold + ) + prepared_input.topk_return_results = ( + topk_return_results or huge_settings.topk_return_results + ) + prepared_input.rerank_method = rerank_method + prepared_input.near_neighbor_first = near_neighbor_first + prepared_input.custom_related_information = custom_related_information + prepared_input.answer_prompt = answer_prompt or prompt.answer_prompt + prepared_input.schema = huge_settings.graph_name + + prepared_input.data_json = { + "query": query, + "vector_search": vector_search, + "graph_search": graph_search, + "max_graph_items": max_graph_items or huge_settings.max_graph_items, + } + return + + def build_flow(self, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, **kwargs) + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + # Create nodes (do not use GRegion, use registerGElement for all nodes) + only_vector_query_node = VectorQueryNode() + merge_rerank_node = MergeRerankNode() + answer_synthesize_node = AnswerSynthesizeNode() + + # Register nodes and dependencies, keep naming consistent with original + pipeline.registerGElement(only_vector_query_node, set(), "only_vector") + pipeline.registerGElement( + merge_rerank_node, {only_vector_query_node}, "merge_two" + ) + pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "vector") + log.info("RAGVectorOnlyFlow pipeline built successfully") + return pipeline + + def post_deal(self, pipeline=None): + if pipeline is None: + return json.dumps( + {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 + ) + try: + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGVectorOnlyFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } + except Exception as e: + log.error(f"RAGVectorOnlyFlow post processing failed: {e}") + return json.dumps( + {"error": f"Post processing failed: {str(e)}"}, + ensure_ascii=False, + indent=2, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index 3aedbe7f2..5afa1bf8e 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -24,6 +24,11 @@ from hugegraph_llm.flows.get_graph_index_info import GetGraphIndexInfoFlow from hugegraph_llm.flows.build_schema import BuildSchemaFlow from hugegraph_llm.flows.prompt_generate import PromptGenerateFlow +from hugegraph_llm.flows.rag_flow_raw import RAGRawFlow +from hugegraph_llm.flows.rag_flow_vector_only import RAGVectorOnlyFlow +from hugegraph_llm.flows.rag_flow_graph_only import RAGGraphOnlyFlow +from hugegraph_llm.flows.rag_flow_graph_vector import RAGGraphVectorFlow +from hugegraph_llm.state.ai_state import WkFlowInput from hugegraph_llm.utils.log import log from hugegraph_llm.flows.text2gremlin import Text2GremlinFlow @@ -67,6 +72,23 @@ def __init__(self, max_pipeline: int = 10): "manager": GPipelineManager(), "flow": Text2GremlinFlow(), } + # New split rag pipelines + self.pipeline_pool["rag_raw"] = { + "manager": GPipelineManager(), + "flow": RAGRawFlow(), + } + self.pipeline_pool["rag_vector_only"] = { + "manager": GPipelineManager(), + "flow": RAGVectorOnlyFlow(), + } + self.pipeline_pool["rag_graph_only"] = { + "manager": GPipelineManager(), + "flow": RAGGraphOnlyFlow(), + } + self.pipeline_pool["rag_graph_vector"] = { + "manager": GPipelineManager(), + "flow": RAGGraphVectorFlow(), + } self.max_pipeline = max_pipeline # TODO: Implement Agentic Workflow @@ -108,6 +130,47 @@ def schedule_flow(self, flow: str, *args, **kwargs): manager.release(pipeline) return res + async def schedule_stream_flow(self, flow: str, *args, **kwargs): + if flow not in self.pipeline_pool: + raise ValueError(f"Unsupported workflow {flow}") + manager: GPipelineManager = self.pipeline_pool[flow]["manager"] + flow: BaseFlow = self.pipeline_pool[flow]["flow"] + pipeline: GPipeline = manager.fetch() + if pipeline is None: + # call coresponding flow_func to create new workflow + pipeline = flow.build_flow(*args, **kwargs) + try: + pipeline.getGParamWithNoEmpty("wkflow_input").stream = True + status = pipeline.init() + if status.isErr(): + error_msg = f"Error in flow init: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + status = pipeline.run() + if status.isErr(): + error_msg = f"Error in flow execution: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + async for res in flow.post_deal_stream(pipeline): + yield res + finally: + manager.add(pipeline) + else: + try: + # fetch pipeline & prepare input for flow + prepared_input: WkFlowInput = pipeline.getGParamWithNoEmpty( + "wkflow_input" + ) + prepared_input.stream = True + flow.prepare(prepared_input, *args, **kwargs) + status = pipeline.run() + if status.isErr(): + raise RuntimeError(f"Error in flow execution {status.getInfo()}") + async for res in flow.post_deal_stream(pipeline): + yield res + finally: + manager.release(pipeline) + class SchedulerSingleton: _instance = None diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py index 0ea0675c0..f90167305 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py @@ -30,6 +30,9 @@ def node_init(self): Node initialization method, can be overridden by subclasses. Returns a CStatus object indicating whether initialization succeeded. """ + if self.wk_input.data_json is not None: + self.context.assign_from_json(self.wk_input.data_json) + self.wk_input.data_json = None return CStatus() def run(self): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py new file mode 100644 index 000000000..78f53e231 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py @@ -0,0 +1,83 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict, Any +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.common_op.merge_dedup_rerank import MergeDedupRerank +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.utils.log import log + + +class MergeRerankNode(BaseNode): + """ + Merge and rerank node, responsible for merging vector and graph query results, deduplication and reranking. + """ + + operator: MergeDedupRerank + + def node_init(self): + """ + Initialize the merge and rerank operator. + """ + try: + # Read user configuration parameters from wk_input + embedding = get_embedding(llm_settings) + graph_ratio = self.wk_input.graph_ratio or 0.5 + rerank_method = self.wk_input.rerank_method or "bleu" + near_neighbor_first = self.wk_input.near_neighbor_first or False + custom_related_information = self.wk_input.custom_related_information or "" + topk_return_results = ( + self.wk_input.topk_return_results or huge_settings.topk_return_results + ) + + self.operator = MergeDedupRerank( + embedding=embedding, + graph_ratio=graph_ratio, + method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + topk_return_results=topk_return_results, + ) + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize MergeRerankNode: {e}") + from PyCGraph import CStatus + + return CStatus(-1, f"MergeRerankNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the merge and rerank operation. + """ + try: + # Perform merge, deduplication, and rerank + result = self.operator.run(data_json) + + # Log result statistics + vector_count = len(result.get("vector_result", [])) + graph_count = len(result.get("graph_result", [])) + merged_count = len(result.get("merged_result", [])) + + log.info( + f"Merge and rerank completed: {vector_count} vector results, " + f"{graph_count} graph results, {merged_count} merged results" + ) + + return result + + except Exception as e: + log.error(f"Merge and rerank failed: {e}") + return data_json 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 4c5acbe97..f71bd7bd5 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 @@ -37,7 +37,7 @@ def node_init(self): if isinstance(texts, str): texts = [texts] self.chunk_split_op = ChunkSplit(texts, split_type, language) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): 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 b576e8170..a4ebc7092 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 @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph import Commit2Graph from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -29,7 +28,7 @@ def node_init(self): if data_json: self.context.assign_from_json(data_json) self.commit_to_graph_op = Commit2Graph() - return CStatus() + return super().node_init() def operator_schedule(self, data_json): return self.commit_to_graph_op.run(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 b2434e524..99b428e5e 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 @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.hugegraph_op.fetch_graph_data import FetchGraphData from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -27,7 +26,7 @@ class FetchGraphDataNode(BaseNode): def node_init(self): self.fetch_graph_data_op = FetchGraphData(get_hg_client()) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): return self.fetch_graph_data_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py new file mode 100644 index 000000000..ae65ccb33 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py @@ -0,0 +1,93 @@ +# 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 PyCGraph import CStatus +from typing import Dict, Any +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery +from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.utils.log import log + + +class GraphQueryNode(BaseNode): + """ + Graph query node, responsible for retrieving relevant information from the graph database. + """ + + graph_rag_query: GraphRAGQuery + + def node_init(self): + """ + Initialize the graph query operator. + """ + try: + graph_name = huge_settings.graph_name + if not graph_name: + return CStatus(-1, "graph_name is required in wk_input") + + max_deep = self.wk_input.max_deep or 2 + max_graph_items = ( + self.wk_input.max_graph_items or huge_settings.max_graph_items + ) + max_v_prop_len = self.wk_input.max_v_prop_len or 2048 + max_e_prop_len = self.wk_input.max_e_prop_len or 256 + prop_to_match = self.wk_input.prop_to_match + num_gremlin_generate_example = self.wk_input.gremlin_tmpl_num or -1 + gremlin_prompt = ( + self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt + ) + + # Initialize GraphRAGQuery operator + self.graph_rag_query = GraphRAGQuery( + max_deep=max_deep, + max_graph_items=max_graph_items, + max_v_prop_len=max_v_prop_len, + max_e_prop_len=max_e_prop_len, + prop_to_match=prop_to_match, + num_gremlin_generate_example=num_gremlin_generate_example, + gremlin_prompt=gremlin_prompt, + ) + + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize GraphQueryNode: {e}") + + return CStatus(-1, f"GraphQueryNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the graph query operation. + """ + try: + # Get the query text from input + query = data_json.get("query", "") + + if not query: + log.warning("No query text provided for graph query") + return data_json + + # Execute the graph query (assuming schema and semantic query have been completed in previous nodes) + graph_result = self.graph_rag_query.run(data_json) + data_json.update(graph_result) + + log.info( + f"Graph query completed, found {len(data_json.get('graph_result', []))} results" + ) + + return data_json + + except Exception as e: + log.error(f"Graph query failed: {e}") + return data_json 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 84719d9eb..3face9d63 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py @@ -15,7 +15,6 @@ import json -from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.common_op.check_schema import CheckSchema from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager @@ -59,7 +58,7 @@ def node_init(self): else: log.info("Get schema '%s' from graphdb.", self.schema) self.schema_manager = self._import_schema(from_hugegraph=self.schema) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): log.debug("SchemaNode input state: %s", 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 ab31fa394..c01cffc91 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 @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from hugegraph_llm.config import llm_settings from hugegraph_llm.models.embeddings.init_embedding import get_embedding from hugegraph_llm.nodes.base_node import BaseNode @@ -28,7 +27,7 @@ class BuildSemanticIndexNode(BaseNode): def node_init(self): self.build_semantic_index_op = BuildSemanticIndex(get_embedding(llm_settings)) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): return self.build_semantic_index_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py index cf2f9b677..1f6a3c75b 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from hugegraph_llm.config import llm_settings from hugegraph_llm.models.embeddings.init_embedding import get_embedding from hugegraph_llm.nodes.base_node import BaseNode @@ -28,7 +27,7 @@ class BuildVectorIndexNode(BaseNode): def node_init(self): self.build_vector_index_op = BuildVectorIndex(get_embedding(llm_settings)) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): return self.build_vector_index_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py index eb033d869..e9283598a 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -19,9 +19,12 @@ from PyCGraph import CStatus +from hugegraph_llm.config import llm_settings from hugegraph_llm.nodes.base_node import BaseNode -from hugegraph_llm.operators.index_op.gremlin_example_index_query import GremlinExampleIndexQuery -from hugegraph_llm.models.embeddings.init_embedding import Embeddings +from hugegraph_llm.operators.index_op.gremlin_example_index_query import ( + GremlinExampleIndexQuery, +) +from hugegraph_llm.models.embeddings.init_embedding import get_embedding class GremlinExampleIndexQueryNode(BaseNode): @@ -29,13 +32,15 @@ class GremlinExampleIndexQueryNode(BaseNode): def node_init(self): # Build operator (index lazy-loading handled in operator) - embedding = Embeddings().get_embedding() + embedding = get_embedding(llm_settings) example_num = getattr(self.wk_input, "example_num", None) if not isinstance(example_num, int): example_num = 2 # Clamp to [0, 10] example_num = max(0, min(10, example_num)) - self.operator = GremlinExampleIndexQuery(embedding=embedding, num_examples=example_num) + self.operator = GremlinExampleIndexQuery( + embedding=embedding, num_examples=example_num + ) return CStatus() def operator_schedule(self, data_json: Dict[str, Any]): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py new file mode 100644 index 000000000..bf605aa49 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -0,0 +1,91 @@ +# 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 PyCGraph import CStatus +from typing import Dict, Any +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.semantic_id_query import SemanticIdQuery +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.utils.log import log + + +class SemanticIdQueryNode(BaseNode): + """ + Semantic ID query node, responsible for semantic matching based on keywords. + """ + + semantic_id_query: SemanticIdQuery + + def node_init(self): + """ + Initialize the semantic ID query operator. + """ + try: + graph_name = huge_settings.graph_name + if not graph_name: + return CStatus(-1, "graph_name is required in wk_input") + + embedding = get_embedding(llm_settings) + by = self.wk_input.semantic_by or "keywords" + topk_per_keyword = ( + self.wk_input.topk_per_keyword or huge_settings.topk_per_keyword + ) + topk_per_query = self.wk_input.topk_per_query or 10 + vector_dis_threshold = ( + self.wk_input.vector_dis_threshold or huge_settings.vector_dis_threshold + ) + + # Initialize the semantic ID query operator + self.semantic_id_query = SemanticIdQuery( + embedding=embedding, + by=by, + topk_per_keyword=topk_per_keyword, + topk_per_query=topk_per_query, + vector_dis_threshold=vector_dis_threshold, + ) + + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize SemanticIdQueryNode: {e}") + + return CStatus(-1, f"SemanticIdQueryNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the semantic ID query operation. + """ + try: + # Get the query text and keywords from input + query = data_json.get("query", "") + keywords = data_json.get("keywords", []) + + if not query and not keywords: + log.warning("No query text or keywords provided for semantic query") + return data_json + + # Perform the semantic query + semantic_result = self.semantic_id_query.run(data_json) + + match_vids = semantic_result.get("match_vids", []) + log.info( + f"Semantic query completed, found {len(match_vids)} matching vertex IDs" + ) + + return semantic_result + + except Exception as e: + log.error(f"Semantic query failed: {e}") + return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py new file mode 100644 index 000000000..48b50acf3 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.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. + +from typing import Dict, Any +from hugegraph_llm.config import llm_settings +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.utils.log import log + + +class VectorQueryNode(BaseNode): + """ + Vector query node, responsible for retrieving relevant documents from the vector index + """ + + operator: VectorIndexQuery + + def node_init(self): + """ + Initialize the vector query operator + """ + try: + # 从 wk_input 中读取用户配置参数 + embedding = get_embedding(llm_settings) + max_items = ( + self.wk_input.max_items if self.wk_input.max_items is not None else 3 + ) + + self.operator = VectorIndexQuery(embedding=embedding, topk=max_items) + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize VectorQueryNode: {e}") + from PyCGraph import CStatus + + return CStatus(-1, f"VectorQueryNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the vector query operation + """ + try: + # Get the query text from input + query = data_json.get("query", "") + if not query: + log.warning("No query text provided for vector query") + return data_json + + # Perform the vector query + result = self.operator.run({"query": query}) + + # Update the state + data_json.update(result) + log.info( + f"Vector query completed, found {len(result.get('vector_result', []))} results" + ) + + return data_json + + except Exception as e: + log.error(f"Vector query failed: {e}") + return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py new file mode 100644 index 000000000..22b970b4a --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py @@ -0,0 +1,99 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict, Any +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize +from hugegraph_llm.utils.log import log + + +class AnswerSynthesizeNode(BaseNode): + """ + Answer synthesis node, responsible for generating the final answer based on retrieval results. + """ + + operator: AnswerSynthesize + + def node_init(self): + """ + Initialize the answer synthesis operator. + """ + try: + prompt_template = self.wk_input.answer_prompt + raw_answer = self.wk_input.raw_answer or False + vector_only_answer = self.wk_input.vector_only_answer or False + graph_only_answer = self.wk_input.graph_only_answer or False + graph_vector_answer = self.wk_input.graph_vector_answer or False + + self.operator = AnswerSynthesize( + prompt_template=prompt_template, + raw_answer=raw_answer, + vector_only_answer=vector_only_answer, + graph_only_answer=graph_only_answer, + graph_vector_answer=graph_vector_answer, + ) + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize AnswerSynthesizeNode: {e}") + from PyCGraph import CStatus + + return CStatus(-1, f"AnswerSynthesizeNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the answer synthesis operation. + """ + try: + if self.getGParamWithNoEmpty("wkflow_input").stream: + # Streaming mode: return a generator for streaming output + data_json["stream_generator"] = self.operator.run_streaming(data_json) + return data_json + else: + # Non-streaming mode: execute answer synthesis + result = self.operator.run(data_json) + + # Record the types of answers generated + answer_types = [] + if result.get("raw_answer"): + answer_types.append("raw") + if result.get("vector_only_answer"): + answer_types.append("vector_only") + if result.get("graph_only_answer"): + answer_types.append("graph_only") + if result.get("graph_vector_answer"): + answer_types.append("graph_vector") + + log.info( + f"Answer synthesis completed for types: {', '.join(answer_types)}" + ) + + # Print enabled answer types according to self.wk_input configuration + wk_input_types = [] + if getattr(self.wk_input, "raw_answer", False): + wk_input_types.append("raw") + if getattr(self.wk_input, "vector_only_answer", False): + wk_input_types.append("vector_only") + if getattr(self.wk_input, "graph_only_answer", False): + wk_input_types.append("graph_only") + if getattr(self.wk_input, "graph_vector_answer", False): + wk_input_types.append("graph_vector") + log.info( + f"Enabled answer types according to wk_input config: {', '.join(wk_input_types)}" + ) + return result + + except Exception as e: + log.error(f"Answer synthesis failed: {e}") + return 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 8bceed804..628765f58 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 @@ -43,7 +43,7 @@ def node_init(self): self.property_graph_extract = PropertyGraphExtract(llm, example_prompt) else: return CStatus(-1, f"Unsupported extract_type: {extract_type}") - return CStatus() + return super().node_init() def operator_schedule(self, data_json): if self.extract_type == "triples": diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py new file mode 100644 index 000000000..76fc06eb3 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py @@ -0,0 +1,80 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict, Any +from PyCGraph import CStatus + +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.llm_op.keyword_extract import KeywordExtract +from hugegraph_llm.utils.log import log + + +class KeywordExtractNode(BaseNode): + operator: KeywordExtract + + """ + Keyword extraction node, responsible for extracting keywords from query text. + """ + + def node_init(self): + """ + Initialize the keyword extraction operator. + """ + try: + max_keywords = ( + self.wk_input.max_keywords + if self.wk_input.max_keywords is not None + else 5 + ) + language = ( + self.wk_input.language + if self.wk_input.language is not None + else "english" + ) + extract_template = self.wk_input.keywords_extract_prompt + + self.operator = KeywordExtract( + text=self.wk_input.query, + max_keywords=max_keywords, + language=language, + extract_template=extract_template, + ) + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize KeywordExtractNode: {e}") + return CStatus(-1, f"KeywordExtractNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the keyword extraction operation. + """ + try: + # Perform keyword extraction + result = self.operator.run(data_json) + if "keywords" not in result: + log.warning("Keyword extraction result missing 'keywords' field") + result["keywords"] = [] + + log.info(f"Extracted keywords: {result.get('keywords', [])}") + + return result + + except Exception as e: + log.error(f"Keyword extraction failed: {e}") + # Add error flag to indicate failure + error_result = data_json.copy() + error_result["error"] = str(e) + error_result["keywords"] = [] + return error_result diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py index 317f9e6ac..8c49994fd 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py @@ -50,7 +50,7 @@ def node_init(self): "example_name": self.wk_input.example_name, } self.context.assign_from_json(context) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): """ diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py index 7df2e68e7..408adb10a 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py @@ -75,7 +75,7 @@ def node_init(self): } self.context.assign_from_json(_context_payload) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): try: diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py index ffbafbaf4..a36831526 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py @@ -22,13 +22,15 @@ from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.config import prompt as prompt_cfg +from hugegraph_llm.models.llms.init_llm import get_text2gql_llm +from hugegraph_llm.config import llm_settings, prompt as prompt_cfg def _stable_schema_string(state_json: Dict[str, Any]) -> str: if "simple_schema" in state_json and state_json["simple_schema"] is not None: - return json.dumps(state_json["simple_schema"], ensure_ascii=False, sort_keys=True) + return json.dumps( + state_json["simple_schema"], ensure_ascii=False, sort_keys=True + ) if "schema" in state_json and state_json["schema"] is not None: return json.dumps(state_json["schema"], ensure_ascii=False, sort_keys=True) return "" @@ -39,7 +41,7 @@ class Text2GremlinNode(BaseNode): def node_init(self): # Select LLM - llm = LLMs().get_text2gql_llm() + llm = get_text2gql_llm(llm_settings) # Serialize schema deterministically state_json = self.context.to_json() schema_str = _stable_schema_string(state_json) diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index f941098b1..3a6fd3c1c 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -24,7 +24,6 @@ class WkFlowInput(GParam): split_type: str = None # split type used by ChunkSplit Node example_prompt: str = None # need by graph information extract schema: str = None # Schema information requeired by SchemaNode - graph_name: str = None data_json = None extract_type = None query_examples = None @@ -34,11 +33,45 @@ class WkFlowInput(GParam): scenario: str = None # Scenario description example_name: str = None # Example name # Fields for Text2Gremlin - query: str = None example_num: int = None gremlin_prompt: str = None requested_outputs: Optional[List[str]] = None + # RAG Flow related fields + query: str = None # User query for RAG + vector_search: bool = None # Enable vector search + graph_search: bool = None # Enable graph search + raw_answer: bool = None # Return raw answer + vector_only_answer: bool = None # Vector only answer mode + graph_only_answer: bool = None # Graph only answer mode + graph_vector_answer: bool = None # Combined graph and vector answer + graph_ratio: float = None # Graph ratio for merging + rerank_method: str = None # Reranking method + near_neighbor_first: bool = None # Near neighbor first flag + custom_related_information: str = None # Custom related information + answer_prompt: str = None # Answer generation prompt + keywords_extract_prompt: str = None # Keywords extraction prompt + gremlin_tmpl_num: int = None # Gremlin template number + gremlin_prompt: str = None # Gremlin generation prompt + max_graph_items: int = None # Maximum graph items + topk_return_results: int = None # Top-k return results + vector_dis_threshold: float = None # Vector distance threshold + topk_per_keyword: int = None # Top-k per keyword + max_keywords: int = None + max_items: int = None + + # Semantic query related fields + semantic_by: str = None # Semantic query method + topk_per_query: int = None # Top-k per query + + # Graph query related fields + max_deep: int = None # Maximum depth for graph traversal + max_v_prop_len: int = None # Maximum vertex property length + max_e_prop_len: int = None # Maximum edge property length + prop_to_match: str = None # Property to match + + stream: bool = None # used for recognize stream mode + def reset(self, _: CStatus) -> None: self.texts = None self.language = None @@ -55,10 +88,40 @@ def reset(self, _: CStatus) -> None: self.scenario = None self.example_name = None # Text2Gremlin related configuration - self.query = None self.example_num = None self.gremlin_prompt = None self.requested_outputs = None + # RAG Flow related fields + self.query = None + self.vector_search = None + self.graph_search = None + self.raw_answer = None + self.vector_only_answer = None + self.graph_only_answer = None + self.graph_vector_answer = None + self.graph_ratio = None + self.rerank_method = None + self.near_neighbor_first = None + self.custom_related_information = None + self.answer_prompt = None + self.keywords_extract_prompt = None + self.gremlin_tmpl_num = None + self.gremlin_prompt = None + self.max_graph_items = None + self.topk_return_results = None + self.vector_dis_threshold = None + self.topk_per_keyword = None + self.max_keywords = None + self.max_items = None + # Semantic query related fields + self.semantic_by = None + self.topk_per_query = None + # Graph query related fields + self.max_deep = None + self.max_v_prop_len = None + self.max_e_prop_len = None + self.prop_to_match = None + self.stream = None class WkFlowState(GParam): @@ -83,6 +146,17 @@ class WkFlowState(GParam): template_exec_res: Optional[Any] = None raw_exec_res: Optional[Any] = None + match_vids = None + vector_result = None + graph_result = None + + raw_answer: str = None + vector_only_answer: str = None + graph_only_answer: str = None + graph_vector_answer: str = None + + merged_result = None + def setup(self): self.schema = None self.simple_schema = None @@ -90,7 +164,7 @@ def setup(self): self.edges = None self.vertices = None self.triples = None - self.call_count = 0 + self.call_count = None self.keywords = None self.vector_result = None @@ -99,12 +173,20 @@ def setup(self): self.generated_extract_prompt = None # Text2Gremlin results reset - self.match_result = [] - self.result = "" - self.raw_result = "" - self.template_exec_res = "" - self.raw_exec_res = "" + self.match_result = None + self.result = None + self.raw_result = None + self.template_exec_res = None + self.raw_exec_res = None + self.raw_answer = None + self.vector_only_answer = None + self.graph_only_answer = None + self.graph_vector_answer = None + + self.vector_result = None + self.graph_result = None + self.merged_result = None return CStatus() def to_json(self): @@ -116,7 +198,11 @@ def to_json(self): dict: A dictionary containing non-None instance members and their serialized values. """ # Only export instance attributes (excluding methods and class attributes) whose values are not None - return {k: v for k, v in self.__dict__.items() if not k.startswith("_") and v is not None} + return { + k: v + for k, v in self.__dict__.items() + if not k.startswith("_") and v is not None + } # Implement a method that assigns keys from data_json as WkFlowState member variables def assign_from_json(self, data_json: dict): diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 7b870033a..3f527f2fa 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -44,30 +44,17 @@ def get_graph_index_info(): raise gr.Error(str(e)) -def get_graph_index_info_old(): - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - graph_summary_info = builder.fetch_graph_data().run() - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(builder.embedding, "model_name", None) - ) - vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) - graph_summary_info["vid_index"] = { - "embed_dim": vector_index.index.d, - "num_vectors": vector_index.index.ntotal, - "num_vids": len(vector_index.properties), - } - return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) - - def clean_all_graph_index(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None), ) - VectorIndex.clean(str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix) + VectorIndex.clean( + str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix + ) VectorIndex.clean( str(os.path.join(resource_path, folder_name, "gremlin_examples")), filename_prefix, @@ -99,14 +86,18 @@ def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: def extract_graph_origin(input_file, input_text, schema, example_prompt) -> str: texts = read_documents(input_file, input_text) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) if not schema: return "ERROR: please input with correct schema/format." error_message = parse_schema(schema, builder) if error_message: return error_message - builder.chunk_split(texts, "document", "zh").extract_info(example_prompt, "property_graph") + builder.chunk_split(texts, "document", "zh").extract_info( + example_prompt, "property_graph" + ) try: context = builder.run() @@ -155,20 +146,6 @@ def update_vid_embedding(): raise gr.Error(str(e)) -def update_vid_embedding_old(): - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - builder.fetch_graph_data().build_vertex_id_semantic_index() - log.debug("Operators: %s", builder.operators) - try: - context = builder.run() - removed_num = context["removed_vid_vector_num"] - added_num = context["added_vid_vector_num"] - return f"Removed {removed_num} vectors, added {added_num} vectors." - except Exception as e: # pylint: disable=broad-exception-caught - log.error(e) - raise gr.Error(str(e)) - - def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: scheduler = SchedulerSingleton.get_instance() @@ -181,73 +158,11 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: return data -def import_graph_data_old(data: str, schema: str) -> Union[str, Dict[str, Any]]: - try: - data_json = json.loads(data.strip()) - log.debug("Import graph data: %s", data) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - if schema: - error_message = parse_schema(schema, builder) - if error_message: - return error_message - - context = builder.commit_to_hugegraph().run(data_json) - gr.Info("Import graph data successfully!") - print(context) - return json.dumps(context, ensure_ascii=False, indent=2) - except Exception as e: # pylint: disable=W0718 - log.error(e) - traceback.print_exc() - # Note: can't use gr.Error here - gr.Warning(str(e) + " Please check the graph data format/type carefully.") - return data - - def build_schema(input_text, query_example, few_shot): scheduler = SchedulerSingleton.get_instance() try: - return scheduler.schedule_flow("build_schema", input_text, query_example, few_shot) + return scheduler.schedule_flow( + "build_schema", input_text, query_example, few_shot + ) except (TypeError, ValueError) as e: raise gr.Error(f"Schema generation failed: {e}") - - -def build_schema_old(input_text, query_example, few_shot): - context = { - "raw_texts": [input_text] if input_text else [], - "query_examples": [], - "few_shot_schema": {}, - } - - if few_shot: - try: - context["few_shot_schema"] = json.loads(few_shot) - except json.JSONDecodeError as e: - raise gr.Error(f"Few Shot Schema is not in a valid JSON format: {e}") from e - - if query_example: - try: - parsed_examples = json.loads(query_example) - # Validate and retain the description and gremlin fields - context["query_examples"] = [ - { - "description": ex.get("description", ""), - "gremlin": ex.get("gremlin", ""), - } - for ex in parsed_examples - if isinstance(ex, dict) and "description" in ex and "gremlin" in ex - ] - except json.JSONDecodeError as e: - raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e - - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - try: - schema = builder.build_schema().run(context) - except Exception as e: - log.error("Failed to generate schema: %s", e) - raise gr.Error(f"Schema generation failed: {e}") from e - try: - formatted_schema = json.dumps(schema, ensure_ascii=False, indent=2) - return formatted_schema - except (TypeError, ValueError) as e: - log.error("Failed to format schema: %s", e) - return str(schema) From 591a0d16a5c4b7c002fbbf60fa5ec24fd5620427 Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Tue, 20 May 2025 20:51:57 +0800 Subject: [PATCH 05/71] feat(llm): index curd test passed --- hugegraph-llm/pyproject.toml | 19 +- .../src/hugegraph_llm/config/__init__.py | 3 + .../src/hugegraph_llm/config/index_config.py | 34 ++++ .../indices/vector_index/base.py | 78 +++++++++ .../faiss_vector_store.py} | 33 +++- .../vector_index/milvus_vector_store.py | 165 ++++++++++++++++++ .../vector_index/qdrant_vector_store.py | 128 ++++++++++++++ ...or_index.py => test_faiss_vector_index.py} | 11 +- .../tests/indices/test_milvus_vector_index.py | 82 +++++++++ .../tests/indices/test_qdrant_vector_index.py | 83 +++++++++ 10 files changed, 624 insertions(+), 12 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/config/index_config.py create mode 100644 hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py rename hugegraph-llm/src/hugegraph_llm/indices/{vector_index.py => vector_index/faiss_vector_store.py} (85%) create mode 100644 hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py create mode 100644 hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py rename hugegraph-llm/src/tests/indices/{test_vector_index.py => test_faiss_vector_index.py} (81%) create mode 100644 hugegraph-llm/src/tests/indices/test_milvus_vector_index.py create mode 100644 hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index 09c49ae26..69f2bf5b9 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -22,11 +22,12 @@ description = "A tool for the implementation and research related to large langu authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, ] +maintainers = [ + { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, +] readme = "README.md" license = "Apache-2.0" requires-python = ">=3.10,<3.12" - - dependencies = [ # Common dependencies "decorator", @@ -87,3 +88,17 @@ allow-direct-references = true [tool.uv.sources] hugegraph-python-client = { workspace = true } pycgraph = { git = "https://github.com/ChunelFeng/CGraph.git", subdirectory = "python", rev = "main", marker = "sys_platform == 'linux'" } + +[tool.mypy] +disable_error_code = ["import-untyped"] +check_untyped_defs = true + +[tool.ruff] +line-length = 120 +indent-width = 4 +extend-exclude = [] + +[tool.ruff.format] +quote-style = "preserve" +indent-style = "space" +line-ending = "auto" diff --git a/hugegraph-llm/src/hugegraph_llm/config/__init__.py b/hugegraph-llm/src/hugegraph_llm/config/__init__.py index 5d4f5782d..426ceb949 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/__init__.py +++ b/hugegraph-llm/src/hugegraph_llm/config/__init__.py @@ -20,6 +20,8 @@ import os +from hugegraph_llm.config.index_config import IndexConfig + from .prompt_config import PromptConfig from .hugegraph_config import HugeGraphConfig from .admin_config import AdminConfig @@ -31,6 +33,7 @@ huge_settings = HugeGraphConfig() admin_settings = AdminConfig() +index_settings = IndexConfig() package_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) resource_path = os.path.join(package_path, "resources") diff --git a/hugegraph-llm/src/hugegraph_llm/config/index_config.py b/hugegraph-llm/src/hugegraph_llm/config/index_config.py new file mode 100644 index 000000000..77d5756a5 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/config/index_config.py @@ -0,0 +1,34 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +from typing import Optional +from .models import BaseConfig +import os + + +class IndexConfig(BaseConfig): + """LLM settings""" + + qdrant_host: Optional[str] = os.environ.get("QDRANT_HOST", None) + qdrant_port: int = int(os.environ.get("QDRANT_PORT", "6333")) + qdrant_api_key: Optional[str] = os.environ.get("QDRANT_API_KEY") if os.environ.get("QDRANT_API_KEY") else None + + milvus_host: Optional[str] = os.environ.get("MILVUS_HOST", None) + milvus_port: int = int(os.environ.get("MILVUS_PORT", "19530")) + milvus_user: str = os.environ.get("MILVUS_USER", "") + milvus_password: str = os.environ.get("MILVUS_PASSWORD", "") diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py new file mode 100644 index 000000000..96978f978 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py @@ -0,0 +1,78 @@ +from abc import ABC, abstractmethod +from typing import List, Any, Union, Set + + +class VectorStoreBase(ABC): + """ + Abstract base class defining the interface for a vector store. + Implementations must support adding, removing, searching vectors, + saving/loading from disk, and cleaning up resources. + """ + + @abstractmethod + def add(self, vectors: List[List[float]], props: List[Any]): + """ + Add a list of vectors and their corresponding properties to the store. + + Args: + vectors (List[List[float]]): List of embedding vectors. + props (List[Any]): List of associated metadata or properties for each vector. + """ + + @abstractmethod + def remove(self, props: Union[Set[Any], List[Any]]) -> int: + """ + Remove vectors based on their associated properties. + + Args: + props (Union[Set[Any], List[Any]]): Properties of vectors to remove. + + Returns: + int: Number of vectors removed. + """ + + @abstractmethod + def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9) -> List[Any]: + """ + Search for the top_k most similar vectors to the query vector. + + Args: + query_vector (List[float]): The vector to query against the index. + top_k (int): Number of top results to return. + dis_threshold (float): Distance threshold below which results are considered relevant. + + Returns: + List[Any]: List of properties of the matched vectors. + """ + + @abstractmethod + def to_index_file(self, dir_path: str): + """ + Persist the vector store (index and metadata) to the specified directory. + + Args: + dir_path (str): Path to the directory where the index and properties will be saved. + """ + + @staticmethod + @abstractmethod + def from_name(dir_path: str) -> "VectorStoreBase": + """ + Load a vector store from the specified directory. + + Args: + dir_path (str): Path to the directory containing the index and properties. + + Returns: + VectorStore: An instance of the vector store. + """ + + @staticmethod + @abstractmethod + def clean(dir_path: str): + """ + Delete the persisted index and properties from the specified directory. + + Args: + dir_path (str): Path to the directory to clean. + """ diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py similarity index 85% rename from hugegraph-llm/src/hugegraph_llm/indices/vector_index.py rename to hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py index f85483185..8135b5bad 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py @@ -23,18 +23,17 @@ import faiss import numpy as np +from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.utils.log import log INDEX_FILE_NAME = "index.faiss" PROPERTIES_FILE_NAME = "properties.pkl" -class VectorIndex: - """Comment""" - +class FaissVectorIndex(VectorStoreBase): def __init__(self, embed_dim: int = 1024): self.index = faiss.IndexFlatL2(embed_dim) - self.properties = [] + self.properties: list[Any] = [] @staticmethod def from_index_file( @@ -87,6 +86,9 @@ def from_index_file( vector_index.index = faiss_index vector_index.properties = properties return vector_index +======= + self.properties: list[Any] = [] +>>>>>>> 9e8cbf9 (feat(llm): index curd test passed):hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py def to_index_file(self, dir_path: str, filename_prefix: str = None): """Save index to files, supporting model-specific filenames.""" @@ -138,12 +140,14 @@ def search( distances, indices = self.index.search(np.array([query_vector]), top_k) results = [] for dist, i in zip(distances[0], indices[0]): - if dist < dis_threshold: # Smaller distances indicate higher similarity + if dist < dis_threshold: results.append(deepcopy(self.properties[i])) log.debug("[✓] Add valid distance %s to results.", dist) else: log.debug( - "[x] Distance %s >= threshold %s, ignore this result.", dist, dis_threshold + "[x] Distance %s >= threshold %s, ignore this result.", + dist, + dis_threshold, ) return results @@ -168,3 +172,20 @@ def clean(dir_path: str, filename_prefix: str = None): log.info("Removed index file: %s", file) except OSError as e: log.error("Error removing file %s: %s", file, e) + + @staticmethod + def from_name(name: str) -> "FaissVectorIndex": + index_file = os.path.join(name, INDEX_FILE_NAME) + properties_file = os.path.join(name, PROPERTIES_FILE_NAME) + if not os.path.exists(index_file) or not os.path.exists(properties_file): + log.warning("No index file found, create a new one.") + return FaissVectorIndex() + + faiss_index = faiss.read_index(index_file) + embed_dim = faiss_index.d + with open(properties_file, "rb") as f: + properties = pkl.load(f) + vector_index = FaissVectorIndex(embed_dim) + vector_index.index = faiss_index + vector_index.properties = properties + return vector_index diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py new file mode 100644 index 000000000..b391aa339 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -0,0 +1,165 @@ +import json +from typing import List, Any, Set, Union + +from pymilvus import ( + connections, + utility, + Collection, + FieldSchema, + CollectionSchema, + DataType, +) + +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +from hugegraph_llm.utils.log import log +from hugegraph_llm.config import index_settings + +COLLECTION_NAME_PREFIX = "hugegraph_llm_" + + +class MilvusVectorIndex(VectorStoreBase): + def __init__( + self, + name: str, + host: str, + port: int, + user="", + password="", + embed_dim: int = 1024, + ): + self.embed_dim = embed_dim + self.host = host + self.port = port + self.name = COLLECTION_NAME_PREFIX + name + connections.connect(host=host, port=port, user=user, password=password) + + if not utility.has_collection(self.name): + self._create_collection() + + self.collection = Collection(self.name) + + def _create_collection(self): + """Create a new collection in Milvus.""" + id_field = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True) + vector_field = FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embed_dim) + property_field = FieldSchema(name="property", dtype=DataType.VARCHAR, max_length=65535) + original_id_field = FieldSchema(name="original_id", dtype=DataType.INT64) + + schema = CollectionSchema( + fields=[id_field, vector_field, property_field, original_id_field], + description="Vector index collection", + ) + + collection = Collection(name=self.name, schema=schema) + + index_params = { + "metric_type": "L2", + "index_type": "IVF_FLAT", + "params": {"nlist": 128}, + } + collection.create_index(field_name="embedding", index_params=index_params) + + def to_index_file(self, name: str): + self.collection.flush() + + def _deserialize_property(self, prop_str): + """Deserialize property from JSON string.""" + try: + return json.loads(prop_str) + except (json.JSONDecodeError, TypeError): + return prop_str + + def add(self, vectors: List[List[float]], props: List[Any]): + if len(vectors) == 0: + return + + # Get the current count to use as starting index + count = self.collection.num_entities + entities = [] + + for i, (vector, prop) in enumerate(zip(vectors, props)): + idx = count + i + entities.append( + { + "embedding": vector, + "property": prop, + "original_id": idx, + } + ) + + self.collection.insert(entities) + self.collection.flush() + + def remove(self, props: Union[Set[Any], List[Any]]) -> int: + if isinstance(props, list): + props = set(props) + try: + self.collection.load() + remove_num = 0 + for prop in props: + expr = f'property == "{prop}"' + res = self.collection.delete(expr) + if hasattr(res, "delete_count"): + remove_num += res.delete_count + if remove_num > 0: + self.collection.flush() + return remove_num + finally: + self.collection.release() + + def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9) -> List[Any]: + try: + if self.collection.num_entities == 0: + return [] + + self.collection.load() + search_params = {"metric_type": "L2", "params": {"nprobe": 10}} + results = self.collection.search( + data=[query_vector], + anns_field="embedding", + param=search_params, + limit=top_k, + output_fields=["property"], + ) + + ret = [] + for hits in results: + for hit in hits: + if hit.distance < dis_threshold: + prop_str = hit.entity.get("property") + prop = self._deserialize_property(prop_str) + ret.append(prop) + log.debug("[✓] Add valid distance %s to results.", hit.distance) + else: + log.debug( + "[x] Distance %s >= threshold %s, ignore this result.", + hit.distance, + dis_threshold, + ) + + return ret + + finally: + self.collection.release() + + @staticmethod + def clean(name: str): + connections.connect( + host=index_settings.milvus_host, + port=index_settings.milvus_port, + user=index_settings.milvus_user, + password=index_settings.milvus_password, + ) + if utility.has_collection(COLLECTION_NAME_PREFIX + name): + utility.drop_collection(COLLECTION_NAME_PREFIX + name) + + @staticmethod + def from_name(name: str) -> "MilvusVectorIndex": + assert index_settings.milvus_host, "Qdrant host is not configured" + return MilvusVectorIndex( + name, + host=index_settings.milvus_host, + port=index_settings.milvus_port, + user=index_settings.milvus_user, + password=index_settings.milvus_password, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py new file mode 100644 index 000000000..13c591ffc --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -0,0 +1,128 @@ +from typing import List, Any, Set, Union + +from qdrant_client import QdrantClient +from qdrant_client.http import models + +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +from hugegraph_llm.utils.log import log +from hugegraph_llm.config import index_settings + +COLLECTION_NAME_PREFIX = "hugegraph_llm_" + + +class QdrantVectorIndex(VectorStoreBase): + def __init__(self, name: str, host: str, port: int, api_key=None, embed_dim: int = 1024): + self.embed_dim = embed_dim + self.host = host + self.port = port + self.name = COLLECTION_NAME_PREFIX + name + self.client = QdrantClient(host=host, port=port, api_key=api_key) + collections = self.client.get_collections().collections + collection_names = [collection.name for collection in collections] + if self.name not in collection_names: + self._create_collection() + + def _create_collection(self): + """Create a new collection in Qdrant.""" + self.client.create_collection( + collection_name=self.name, + vectors_config=models.VectorParams(size=self.embed_dim, distance=models.Distance.COSINE), + ) + log.info("Created Qdrant collection '%s'", self.name) + + def to_index_file(self, name: str): + # nothing to do when qdrant + pass + + def add(self, vectors: List[List[float]], props: List[Any]): + if len(vectors) == 0: + return + + points = [] + + for i, (vector, prop) in enumerate(zip(vectors, props)): + points.append( + models.PointStruct( + id=i, + vector=vector, + payload={"property": prop}, + ) + ) + + self.client.upsert(collection_name=self.name, points=points, wait=True) + + def remove(self, props: Union[Set[Any], List[Any]]) -> int: + if isinstance(props, list): + props = set(props) + + remove_num = 0 + + for prop in props: + serialized_prop = prop + search_result = self.client.scroll( + collection_name=self.name, + scroll_filter=models.Filter( + must=[ + models.FieldCondition( + key="property", + match=models.MatchValue(value=serialized_prop), + ) + ] + ), + limit=1000, + ) + if search_result and search_result[0]: + point_ids = [point.id for point in search_result[0]] + + if point_ids: + _ = self.client.delete( + collection_name=self.name, + points_selector=models.PointIdsList(points=point_ids), + wait=True, + ) + remove_num += len(point_ids) + + return remove_num + + def search(self, query_vector: List[float], top_k: int = 5, dis_threshold: float = 0.9): + search_result = self.client.search(collection_name=self.name, query_vector=query_vector, limit=top_k) + + result_properties = [] + + for hit in search_result: + distance = 1.0 - hit.score + if distance < dis_threshold: + if hit.payload is not None: + result_properties.append(hit.payload.get("property")) + log.debug("[✓] Add valid distance %s to results.", distance) + else: + log.debug("[x] Hit payload is None, skipping.") + else: + log.debug( + "[x] Distance %s >= threshold %s, ignore this result.", + distance, + dis_threshold, + ) + + return result_properties + + @staticmethod + def clean(name: str): + client = QdrantClient( + host=index_settings.qdrant_host, port=index_settings.qdrant_port, api_key=index_settings.qdrant_api_key + ) + collections = client.get_collections().collections + collection_names = [collection.name for collection in collections] + name = COLLECTION_NAME_PREFIX + name + if name in collection_names: + client.delete_collection(collection_name=name) + + @staticmethod + def from_name(name: str) -> "QdrantVectorIndex": + assert index_settings.qdrant_host, "Qdrant host is not configured" + return QdrantVectorIndex( + name=COLLECTION_NAME_PREFIX + name, + host=index_settings.qdrant_host, + port=index_settings.qdrant_port, + api_key=index_settings.qdrant_api_key, + ) diff --git a/hugegraph-llm/src/tests/indices/test_vector_index.py b/hugegraph-llm/src/tests/indices/test_faiss_vector_index.py similarity index 81% rename from hugegraph-llm/src/tests/indices/test_vector_index.py rename to hugegraph-llm/src/tests/indices/test_faiss_vector_index.py index 0f8fd5f48..0ea3ed316 100644 --- a/hugegraph-llm/src/tests/indices/test_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_faiss_vector_index.py @@ -18,17 +18,20 @@ import unittest from pprint import pprint - -from hugegraph_llm.indices.vector_index import VectorIndex from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex class TestVectorIndex(unittest.TestCase): def test_vector_index(self): embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") - data = ["腾讯的合伙人有字节跳动", "谷歌和微软是竞争关系", "美团的合伙人有字节跳动"] + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] data_embedding = [embedder.get_text_embedding(d) for d in data] - index = VectorIndex(1024) + index = FaissVectorIndex(1024) index.add(data_embedding, data) query = "腾讯的合伙人有哪些？" query_vector = embedder.get_text_embedding(query) diff --git a/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py new file mode 100644 index 000000000..8288b4563 --- /dev/null +++ b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py @@ -0,0 +1,82 @@ +import unittest +from pprint import pprint + +from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding +from hugegraph_llm.indices.vector_index.milvus_vector_store import MilvusVectorIndex + +test_name = "test" + + +class TestMilvusVectorIndex(unittest.TestCase): + def tearDown(self): + MilvusVectorIndex.clean(test_name) + + def test_vector_index(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = MilvusVectorIndex.from_name(test_name) + index.add(data_embedding, data) + + query = "腾讯的合伙人有哪些？" + query_vector = embedder.get_text_embedding(query) + results = index.search(query_vector, 2, dis_threshold=1000) + pprint(results) + + self.assertIsNotNone(results) + self.assertLessEqual(len(results), 2) + + def test_save_and_load(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = MilvusVectorIndex.from_name(test_name) + index.add(data_embedding, data) + + index.to_index_file(test_name) + + loaded_index = MilvusVectorIndex.from_name(test_name) + + query = "腾讯的合伙人有哪些？" + query_vector = embedder.get_text_embedding(query) + results = loaded_index.search(query_vector, 2, dis_threshold=1000) + + self.assertIsNotNone(results) + self.assertLessEqual(len(results), 2) + + def test_remove_entries(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = MilvusVectorIndex.from_name(test_name) + index.add(data_embedding, data) + + query = "合伙人" + query_vector = embedder.get_text_embedding(query) + initial_results = index.search(query_vector, 3, dis_threshold=1000) + initial_count = len(initial_results) + + remove_count = index.remove(["谷歌和微软是竞争关系"]) + + self.assertEqual(remove_count, 1) + + after_results = index.search(query_vector, 3, dis_threshold=1000) + self.assertLessEqual(len(after_results), initial_count - 1) + self.assertNotIn("谷歌和微软是竞争关系", after_results) diff --git a/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py new file mode 100644 index 000000000..f879b9cdd --- /dev/null +++ b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py @@ -0,0 +1,83 @@ +import unittest +from pprint import pprint +from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding +from hugegraph_llm.indices.vector_index.qdrant_vector_store import QdrantVectorIndex + + +class TestQdrantVectorIndex(unittest.TestCase): + def setUp(self): + self.name = "test" + + def tearDown(self): + QdrantVectorIndex.clean(self.name) + + def test_vector_index(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = QdrantVectorIndex.from_name(self.name) + index.add(data_embedding, data) + + query = "腾讯的合伙人有哪些？" + query_vector = embedder.get_text_embedding(query) + results = index.search(query_vector, 2, dis_threshold=100) + pprint(results) + + self.assertIsNotNone(results) + self.assertLessEqual(len(results), 2) + + def test_save_and_load(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = QdrantVectorIndex.from_name(self.name) + index.add(data_embedding, data) + + index.to_index_file(self.name) + + loaded_index = QdrantVectorIndex.from_name(self.name) + + query = "腾讯的合伙人有哪些？" + query_vector = embedder.get_text_embedding(query) + results = loaded_index.search(query_vector, 2, dis_threshold=100) + + self.assertIsNotNone(results) + self.assertLessEqual(len(results), 2) + + def test_remove_entries(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = QdrantVectorIndex.from_name(self.name) + index.add(data_embedding, data) + + query = "合伙人" + query_vector = embedder.get_text_embedding(query) + initial_results = index.search(query_vector, 3, dis_threshold=100) + initial_count = len(initial_results) + + remove_count = index.remove(["谷歌和微软是竞争关系"]) + + self.assertEqual(remove_count, 1) + + after_results = index.search(query_vector, 3) + self.assertLessEqual(len(after_results), initial_count - 1) + self.assertNotIn("谷歌和微软是竞争关系", after_results) From 0692bcabf9b942d625c54b96164a46e7331e9e88 Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Tue, 20 May 2025 21:14:19 +0800 Subject: [PATCH 06/71] feat(llm): some type bug && revert to FaissVectorIndex --- hugegraph-llm/pyproject.toml | 1 + .../hugegraph_llm/config/hugegraph_config.py | 22 ++++---- .../vector_index/faiss_vector_store.py | 50 +++++++++++++++++++ .../operators/hugegraph_op/schema_manager.py | 2 +- .../index_op/build_gremlin_example_index.py | 4 +- .../index_op/build_semantic_index.py | 4 +- .../operators/index_op/build_vector_index.py | 11 ++++ .../index_op/gremlin_example_index_query.py | 28 +++++++++-- .../operators/index_op/semantic_id_query.py | 22 ++++++++ .../operators/index_op/vector_index_query.py | 20 ++++++++ .../operators/kg_construction_task.py | 5 ++ .../hugegraph_llm/utils/graph_index_utils.py | 44 +++++++++------- .../hugegraph_llm/utils/vector_index_utils.py | 13 +++-- 13 files changed, 182 insertions(+), 44 deletions(-) diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index 69f2bf5b9..f8227deca 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -92,6 +92,7 @@ pycgraph = { git = "https://github.com/ChunelFeng/CGraph.git", subdirectory = "p [tool.mypy] disable_error_code = ["import-untyped"] check_untyped_defs = true +disallow_untyped_defs = false [tool.ruff] line-length = 120 diff --git a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py b/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py index 69abf0fbc..18710bcdf 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py @@ -23,21 +23,21 @@ class HugeGraphConfig(BaseConfig): """HugeGraph settings""" # graph server config - graph_url: Optional[str] = "127.0.0.1:8080" - graph_name: Optional[str] = "hugegraph" - graph_user: Optional[str] = "admin" - graph_pwd: Optional[str] = "xxx" + graph_url: str = "127.0.0.1:8080" + graph_name: str = "hugegraph" + graph_user: str = "admin" + graph_pwd: str = "xxx" graph_space: Optional[str] = None # graph query config - limit_property: Optional[str] = "False" - max_graph_path: Optional[int] = 10 - max_graph_items: Optional[int] = 30 - edge_limit_pre_label: Optional[int] = 8 + limit_property: str = "False" + max_graph_path: int = 10 + max_graph_items: int = 30 + edge_limit_pre_label: int = 8 # vector config - vector_dis_threshold: Optional[float] = 0.9 - topk_per_keyword: Optional[int] = 1 + vector_dis_threshold: float = 0.9 + topk_per_keyword: int = 1 # rerank config - topk_return_results: Optional[int] = 20 + topk_return_results: int = 20 diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py index 8135b5bad..1eed19492 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py @@ -151,6 +151,56 @@ def search( ) return results + @staticmethod + def from_index_file( + dir_path: str, filename_prefix: str | None = None, record_miss: bool = True + ) -> "FaissVectorIndex": + """Load index from files, supporting model-specific filenames. + + If prefixed files are missing, optionally warn and return an empty index. + Also validates vector/property count consistency. + """ + index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME + property_name = ( + f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + ) + index_file = os.path.join(dir_path, index_name) + properties_file = os.path.join(dir_path, property_name) + missing = [p for p in [index_file, properties_file] if not os.path.exists(p)] + if missing: + if record_miss: + log.warning( + "Missing vector files: %s. Need create a new one for it.", ", ".join(missing) + ) + return FaissVectorIndex() + + try: + faiss_index = faiss.read_index(index_file) + with open(properties_file, "rb") as f: + properties = pkl.load(f) + except (RuntimeError, pkl.UnpicklingError, OSError) as e: # pragma: no cover + log.error( + "Failed to load index files for model '%s': %s", + filename_prefix or "default", + e, + ) + raise RuntimeError( + f"Could not load index files for model '{filename_prefix or 'default'}'. " + f"Original error ({type(e).__name__}): {e}" + ) from e + + if faiss_index.ntotal != len(properties): + raise RuntimeError( + f"Data inconsistency: index for model '{filename_prefix or 'default'}' has " + f"{faiss_index.ntotal} vectors, but {len(properties)} properties." + ) + + embed_dim = faiss_index.d + vector_index = FaissVectorIndex(embed_dim) + vector_index.index = faiss_index + vector_index.properties = properties + return vector_index + @staticmethod def clean(dir_path: str, filename_prefix: str = None): """Clean index files, supporting model-specific filenames. 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 90f1c00ea..2f0643a77 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 @@ -33,7 +33,7 @@ def __init__(self, graph_name: str): self.schema = self.client.schema() def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: - mini_schema = {} + mini_schema = {} # type: ignore # Add necessary vertexlabels items (3) if "vertexlabels" in schema: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index 6d9f96214..fb5c16acb 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -21,7 +21,7 @@ from typing import Dict, Any, List from hugegraph_llm.config import resource_path, llm_settings, huge_settings -from hugegraph_llm.indices.vector_index import VectorIndex +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.utils.embedding_utils import ( get_embeddings_parallel, @@ -50,7 +50,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: examples_embedding = asyncio.run(get_embeddings_parallel(self.embedding, queries)) embed_dim = len(examples_embedding[0]) if len(self.examples) > 0: - vector_index = VectorIndex(embed_dim) + vector_index = FaissVectorIndex(embed_dim) vector_index.add(examples_embedding, self.examples) vector_index.to_index_file(self.index_dir, self.filename_prefix) context["embed_dim"] = embed_dim 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 5689a59ac..ee6837f96 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 @@ -21,7 +21,7 @@ from typing import Any, Dict from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index import VectorIndex +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager from hugegraph_llm.utils.embedding_utils import ( @@ -41,7 +41,7 @@ def __init__(self, embedding: BaseEmbedding): self.filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(embedding, "model_name", None) ) - self.vid_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) + self.vid_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) self.embedding = embedding self.sm = SchemaManager(huge_settings.graph_name) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py index f5fb823c5..c9aa1876a 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py @@ -20,8 +20,14 @@ import os from typing import Dict, Any +<<<<<<< HEAD from hugegraph_llm.config import huge_settings, resource_path, llm_settings from hugegraph_llm.indices.vector_index import VectorIndex +======= +from tqdm import tqdm +from hugegraph_llm.config import huge_settings, resource_path +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.utils.embedding_utils import ( get_embeddings_parallel, @@ -34,6 +40,7 @@ class BuildVectorIndex: def __init__(self, embedding: BaseEmbedding): self.embedding = embedding +<<<<<<< HEAD self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space ) @@ -42,6 +49,10 @@ def __init__(self, embedding: BaseEmbedding): llm_settings.embedding_type, getattr(self.embedding, "model_name", None) ) self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) +======= + self.index_dir = str(os.path.join(resource_path, huge_settings.graph_name, "chunks")) + self.vector_index = FaissVectorIndex.from_name(self.index_dir) +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: if "chunks" not in context: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index b680f2ca3..89861ead2 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -18,12 +18,17 @@ import asyncio import os -from typing import Dict, Any, List +from typing import Dict, Any, List, Optional import pandas as pd +<<<<<<< HEAD from hugegraph_llm.config import resource_path, llm_settings, huge_settings from hugegraph_llm.indices.vector_index import VectorIndex, INDEX_FILE_NAME, PROPERTIES_FILE_NAME +======= +from hugegraph_llm.config import resource_path, huge_settings, llm_settings +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.utils.embedding_utils import ( @@ -35,7 +40,7 @@ class GremlinExampleIndexQuery: - def __init__(self, embedding: BaseEmbedding = None, num_examples: int = 1): + def __init__(self, embedding: Optional[BaseEmbedding] = None, num_examples: int = 1): self.embedding = embedding or Embeddings().get_embedding() self.num_examples = num_examples self.folder_name = get_index_folder_name( @@ -46,7 +51,7 @@ def __init__(self, embedding: BaseEmbedding = None, num_examples: int = 1): llm_settings.embedding_type, getattr(self.embedding, "model_name", None) ) self._ensure_index_exists() - self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) + self.vector_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) def _ensure_index_exists(self): index_name = ( @@ -74,13 +79,28 @@ def _get_match_result(self, context: Dict[str, Any], query: str) -> List[Dict[st return self.vector_index.search(query_embedding, self.num_examples, dis_threshold=1.8) def _build_default_example_index(self): +<<<<<<< HEAD properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict( orient="records" ) # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) queries = [row["query"] for row in properties] embeddings = asyncio.run(get_embeddings_parallel(self.embedding, queries)) - vector_index = VectorIndex(len(embeddings[0])) + vector_index = FaissVectorIndex(len(embeddings[0])) +======= + properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict(orient="records") + from concurrent.futures import ThreadPoolExecutor + + # TODO: use asyncio for IO tasks + with ThreadPoolExecutor() as executor: + embeddings = list( + tqdm( + executor.map(self.embedding.get_text_embedding, [row["query"] for row in properties]), + total=len(properties), + ) + ) + vector_index = FaissVectorIndex(len(embeddings[0])) +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) vector_index.add(embeddings, properties) vector_index.to_index_file(self.index_dir, self.filename_prefix) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py index 3ac03246f..97b26c2c2 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py @@ -19,8 +19,13 @@ import os from typing import Dict, Any, Literal, List, Tuple +<<<<<<< HEAD from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index import VectorIndex +======= +from hugegraph_llm.config import resource_path, huge_settings, llm_settings +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.utils.embedding_utils import get_filename_prefix, get_index_folder_name from hugegraph_llm.utils.log import log @@ -38,6 +43,7 @@ def __init__( topk_per_keyword: int = huge_settings.topk_per_keyword, vector_dis_threshold: float = huge_settings.vector_dis_threshold, ): +<<<<<<< HEAD self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space ) @@ -46,6 +52,16 @@ def __init__( llm_settings.embedding_type, getattr(embedding, "model_name", None) ) self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) +======= + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(embedding, "model_name", None) + ) + self.vector_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) self.embedding = embedding self.by = by self.topk_per_query = topk_per_query @@ -82,11 +98,17 @@ def _exact_match_vids(self, keywords: List[str]) -> Tuple[List[str], List[str]]: def _fuzzy_match_vids(self, keywords: List[str]) -> List[str]: fuzzy_match_result = [] for keyword in keywords: +<<<<<<< HEAD keyword_vector = self.embedding.get_texts_embeddings([keyword])[0] results = self.vector_index.search( keyword_vector, top_k=self.topk_per_keyword, dis_threshold=float(self.vector_dis_threshold), +======= + keyword_vector = self.embedding.get_texts_embeddings([keyword])[0] + results = self.vector_index.search( + keyword_vector, top_k=self.topk_per_keyword, dis_threshold=float(self.vector_dis_threshold) +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) ) if results: fuzzy_match_result.extend(results[: self.topk_per_keyword]) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py index 4ed616929..c34007d1e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py @@ -19,8 +19,13 @@ import os from typing import Dict, Any +<<<<<<< HEAD from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index import VectorIndex +======= +from hugegraph_llm.config import resource_path, huge_settings, llm_settings +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.utils.embedding_utils import get_filename_prefix, get_index_folder_name from hugegraph_llm.utils.log import log @@ -30,6 +35,7 @@ class VectorIndexQuery: def __init__(self, embedding: BaseEmbedding, topk: int = 3): self.embedding = embedding self.topk = topk +<<<<<<< HEAD self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space ) @@ -42,6 +48,20 @@ def __init__(self, embedding: BaseEmbedding, topk: int = 3): def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query = context.get("query") query_embedding = self.embedding.get_texts_embeddings([query])[0] +======= + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(embedding, "model_name", None) + ) + self.vector_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) + + def run(self, context: Dict[str, Any]) -> Dict[str, Any]: + query = context.get("query") + query_embedding = self.embedding.get_texts_embeddings([query])[0] +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) # TODO: why set dis_threshold=2? results = self.vector_index.search(query_embedding, self.topk, dis_threshold=2) # TODO: check format results diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py index 3b5c63103..dac1fa047 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py @@ -37,6 +37,7 @@ class KgBuilder: +<<<<<<< HEAD def __init__( self, llm: BaseLLM, @@ -44,6 +45,10 @@ def __init__( graph: Optional[PyHugeClient] = None, ): self.operators = [] +======= + def __init__(self, llm: BaseLLM, embedding: Optional[BaseEmbedding] = None, graph: Optional[PyHugeClient] = None): + self.operators: List[Any] = [] +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) self.llm = llm self.embedding = embedding self.graph = graph diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 3f527f2fa..9645d893e 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -28,37 +28,43 @@ from .hugegraph_utils import get_hg_client, clean_hg_data from .log import log from .vector_index_utils import read_documents +<<<<<<< HEAD from ..config import resource_path, huge_settings, llm_settings -from ..indices.vector_index import VectorIndex +from ..indices.vector_index.faiss_vector_store import FaissVectorIndex +======= +from ..config import resource_path, huge_settings +from ..indices.vector_index.faiss_vector_store import FaissVectorIndex +>>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) from ..models.embeddings.init_embedding import Embeddings from ..models.llms.init_llm import LLMs from ..operators.kg_construction_task import KgBuilder def get_graph_index_info(): - try: - scheduler = SchedulerSingleton.get_instance() - return scheduler.schedule_flow("get_graph_index_info") - except Exception as e: # pylint: disable=broad-exception-caught - log.error(e) - raise gr.Error(str(e)) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + graph_summary_info = builder.fetch_graph_data().run() + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None) + ) + vector_index = FaissVectorIndex.from_index_file( + str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix, record_miss=False + ) + graph_summary_info["vid_index"] = { + "embed_dim": vector_index.index.d, + "num_vectors": vector_index.index.ntotal, + "num_vids": len(vector_index.properties), + } + return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) def clean_all_graph_index(): - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( - llm_settings.embedding_type, - getattr(Embeddings().get_embedding(), "model_name", None), - ) - VectorIndex.clean( - str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix - ) - VectorIndex.clean( - str(os.path.join(resource_path, folder_name, "gremlin_examples")), - filename_prefix, + llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None) ) + FaissVectorIndex.clean(str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix) + FaissVectorIndex.clean(str(os.path.join(resource_path, folder_name, "gremlin_examples")), filename_prefix) log.warning("Clear graph index and text2gql index successfully!") gr.Info("Clear graph index and text2gql index successfully!") diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 301a6bdab..f7610ce9a 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -22,13 +22,16 @@ import gradio as gr from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index import VectorIndex -from hugegraph_llm.models.embeddings.init_embedding import model_map +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex +from hugegraph_llm.models.embeddings.init_embedding import Embeddings, model_map from hugegraph_llm.flows.scheduler import SchedulerSingleton from hugegraph_llm.utils.embedding_utils import ( get_filename_prefix, get_index_folder_name, ) +from hugegraph_llm.models.llms.init_llm import LLMs +from hugegraph_llm.operators.kg_construction_task import KgBuilder +from hugegraph_llm.utils.hugegraph_utils import get_hg_client def read_documents(input_file, input_text): @@ -64,12 +67,12 @@ def get_vector_index_info(): filename_prefix = get_filename_prefix( llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) ) - chunk_vector_index = VectorIndex.from_index_file( + chunk_vector_index = FaissVectorIndex.from_index_file( str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix, record_miss=False, ) - graph_vid_vector_index = VectorIndex.from_index_file( + graph_vid_vector_index = FaissVectorIndex.from_index_file( str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix ) return json.dumps( @@ -91,7 +94,7 @@ def clean_vector_index(): filename_prefix = get_filename_prefix( llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) ) - VectorIndex.clean(str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix) + FaissVectorIndex.clean(str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix) gr.Info("Clean vector index successfully!") From 70da9936f948e16c161c9dbab495470f04cf67c0 Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Tue, 20 May 2025 21:16:34 +0800 Subject: [PATCH 07/71] feat(llm): some type bug --- .../src/hugegraph_llm/operators/kg_construction_task.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py index dac1fa047..7cea5f75a 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py @@ -86,6 +86,7 @@ def extract_info( if extract_type == "triples": self.operators.append(InfoExtract(self.llm, example_prompt)) elif extract_type == "property_graph": + assert example_prompt self.operators.append(PropertyGraphExtract(self.llm, example_prompt)) return self @@ -98,10 +99,12 @@ def commit_to_hugegraph(self): return self def build_vertex_id_semantic_index(self): + assert self.embedding self.operators.append(BuildSemanticIndex(self.embedding)) return self def build_vector_index(self): + assert self.embedding self.operators.append(BuildVectorIndex(self.embedding)) return self @@ -118,6 +121,7 @@ def build_schema(self): def run(self, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: for operator in self.operators: context = self._run_operator(operator, context) + assert context is not None return context @log_operator_time From 68fd9747a2babe0e5112fc2e01d377fb920f9af0 Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Tue, 20 May 2025 21:28:30 +0800 Subject: [PATCH 08/71] feat(llm): some type bug(from mypy) --- hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 eabd5dd9b..b0979763c 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -115,7 +115,7 @@ def init_rag_ui() -> gr.Interface: def refresh_ui_config_prompt() -> tuple: # we can use its __init__() for in-place reload # settings.from_env() - huge_settings.__init__() # pylint: disable=C2801 + huge_settings.__init__() # type: ignore[misc] # pylint: disable=C2801 prompt.ensure_yaml_file_exists() return ( huge_settings.graph_url, @@ -164,6 +164,7 @@ def create_app(): # we don't need to manually check the env now # settings.check_env() prompt.update_yaml_file() + assert admin_settings.enable_login auth_enabled = admin_settings.enable_login.lower() == "true" log.info("(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled") api_auth = APIRouter(dependencies=[Depends(authenticate)] if auth_enabled else []) From 997d2e21703b64c508e8f395a6d636a47a719aa4 Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Tue, 20 May 2025 21:43:52 +0800 Subject: [PATCH 09/71] feat(llm): add License header --- .../hugegraph_llm/indices/vector_index/base.py | 18 ++++++++++++++++++ .../vector_index/milvus_vector_store.py | 17 +++++++++++++++++ .../vector_index/qdrant_vector_store.py | 17 +++++++++++++++++ .../tests/indices/test_milvus_vector_index.py | 18 ++++++++++++++++++ .../tests/indices/test_qdrant_vector_index.py | 18 ++++++++++++++++++ 5 files changed, 88 insertions(+) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py index 96978f978..ba42289df 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py @@ -1,3 +1,21 @@ +# 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 abc import ABC, abstractmethod from typing import List, Any, Union, Set diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py index b391aa339..503bd1b5f 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -1,3 +1,20 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + import json from typing import List, Any, Set, Union diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index 13c591ffc..f7931b415 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -1,3 +1,20 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + from typing import List, Any, Set, Union from qdrant_client import QdrantClient diff --git a/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py index 8288b4563..b339111a4 100644 --- a/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py @@ -1,3 +1,21 @@ +# 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 unittest from pprint import pprint diff --git a/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py index f879b9cdd..53288d662 100644 --- a/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py @@ -1,3 +1,21 @@ +# 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 unittest from pprint import pprint from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding From 58777ffd9d45eb842e5831e196e5a47963cee5ff Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Tue, 20 May 2025 21:51:36 +0800 Subject: [PATCH 10/71] feat(llm): import sort && change name --- hugegraph-llm/pyproject.toml | 5 ++ .../src/hugegraph_llm/config/index_config.py | 4 +- .../src/hugegraph_llm/config/llm_config.py | 2 +- .../indices/vector_index/base.py | 14 ++--- .../vector_index/faiss_vector_store.py | 58 +------------------ .../vector_index/milvus_vector_store.py | 10 ++-- .../vector_index/qdrant_vector_store.py | 4 +- .../tests/indices/test_faiss_vector_index.py | 3 +- .../tests/indices/test_milvus_vector_index.py | 2 +- .../tests/indices/test_qdrant_vector_index.py | 3 +- 10 files changed, 30 insertions(+), 75 deletions(-) diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index f8227deca..e5e1f502d 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -98,6 +98,11 @@ disallow_untyped_defs = false line-length = 120 indent-width = 4 extend-exclude = [] +extend-select = ["I"] +<<<<<<< HEAD +======= +extend-select = ["I"] +>>>>>>> 56bcac3 (feat(llm): import sort && change name) [tool.ruff.format] quote-style = "preserve" diff --git a/hugegraph-llm/src/hugegraph_llm/config/index_config.py b/hugegraph-llm/src/hugegraph_llm/config/index_config.py index 77d5756a5..b55353476 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/index_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/index_config.py @@ -15,10 +15,10 @@ # specific language governing permissions and limitations # under the License. - +import os from typing import Optional + from .models import BaseConfig -import os class IndexConfig(BaseConfig): diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index 916d70ddf..304e5ac44 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -17,7 +17,7 @@ import os -from typing import Optional, Literal +from typing import Literal, Optional from .models import BaseConfig diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py index ba42289df..53743b2eb 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py @@ -17,7 +17,7 @@ from abc import ABC, abstractmethod -from typing import List, Any, Union, Set +from typing import Any, List, Set, Union class VectorStoreBase(ABC): @@ -64,22 +64,22 @@ def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0 """ @abstractmethod - def to_index_file(self, dir_path: str): + def to_index_file(self, name: str): """ Persist the vector store (index and metadata) to the specified directory. Args: - dir_path (str): Path to the directory where the index and properties will be saved. + name (str): Path to the directory where the index and properties will be saved. """ @staticmethod @abstractmethod - def from_name(dir_path: str) -> "VectorStoreBase": + def from_name(name: str) -> "VectorStoreBase": """ Load a vector store from the specified directory. Args: - dir_path (str): Path to the directory containing the index and properties. + name (str): Path to the directory containing the index and properties. Returns: VectorStore: An instance of the vector store. @@ -87,10 +87,10 @@ def from_name(dir_path: str) -> "VectorStoreBase": @staticmethod @abstractmethod - def clean(dir_path: str): + def clean(name: str): """ Delete the persisted index and properties from the specified directory. Args: - dir_path (str): Path to the directory to clean. + name (str): Path to the directory to clean. """ diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py index 1eed19492..934f4b719 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py @@ -18,7 +18,7 @@ import os import pickle as pkl from copy import deepcopy -from typing import List, Any, Set, Union +from typing import Any, List, Set, Union import faiss import numpy as np @@ -35,60 +35,7 @@ def __init__(self, embed_dim: int = 1024): self.index = faiss.IndexFlatL2(embed_dim) self.properties: list[Any] = [] - @staticmethod - def from_index_file( - dir_path: str, filename_prefix: str = None, record_miss: bool = True - ) -> "VectorIndex": - """Load index from files, supporting model-specific filenames. - - This method loads a Faiss index and its corresponding properties from a directory. - It handles model-specific filenames by constructing them inline using f-strings. - If the specified files are not found, it returns a new, empty VectorIndex instance. - It also performs a consistency check to ensure the number of vectors in the index - matches the number of properties. - """ - index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = ( - f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME - ) - index_file = os.path.join(dir_path, index_name) - properties_file = os.path.join(dir_path, property_name) - miss_files = [f for f in [index_file, properties_file] if not os.path.exists(f)] - if miss_files: - if record_miss: - log.warning( - "Missing vector files: %s. \nNeed create a new one for it.", - ", ".join(miss_files), - ) - return VectorIndex() - - try: - faiss_index = faiss.read_index(index_file) - with open(properties_file, "rb") as f: - properties = pkl.load(f) - except (RuntimeError, pkl.UnpicklingError, OSError) as e: - log.error( - "Failed to load index files for model '%s': %s", filename_prefix or "default", e - ) - raise RuntimeError( - f"Could not load index files for model '{filename_prefix or 'default'}'. " - f"Original error ({type(e).__name__}): {e}" - ) from e - - if faiss_index.ntotal != len(properties): - raise RuntimeError( - f"Data inconsistency: index for model '{filename_prefix or 'default'}' has " - f"{faiss_index.ntotal} vectors, but {len(properties)} properties." - ) - - embed_dim = faiss_index.d - vector_index = VectorIndex(embed_dim) - vector_index.index = faiss_index - vector_index.properties = properties - return vector_index -======= - self.properties: list[Any] = [] ->>>>>>> 9e8cbf9 (feat(llm): index curd test passed):hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py + def to_index_file(self, dir_path: str, filename_prefix: str = None): """Save index to files, supporting model-specific filenames.""" @@ -222,6 +169,7 @@ def clean(dir_path: str, filename_prefix: str = None): log.info("Removed index file: %s", file) except OSError as e: log.error("Error removing file %s: %s", file, e) + @staticmethod def from_name(name: str) -> "FaissVectorIndex": diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py index 503bd1b5f..505209e16 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -16,20 +16,20 @@ # under the License. import json -from typing import List, Any, Set, Union +from typing import Any, List, Set, Union from pymilvus import ( - connections, - utility, Collection, - FieldSchema, CollectionSchema, DataType, + FieldSchema, + connections, + utility, ) +from hugegraph_llm.config import index_settings from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.utils.log import log -from hugegraph_llm.config import index_settings COLLECTION_NAME_PREFIX = "hugegraph_llm_" diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index f7931b415..7f392b184 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -15,14 +15,14 @@ # specific language governing permissions and limitations # under the License. -from typing import List, Any, Set, Union +from typing import Any, List, Set, Union from qdrant_client import QdrantClient from qdrant_client.http import models +from hugegraph_llm.config import index_settings from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.utils.log import log -from hugegraph_llm.config import index_settings COLLECTION_NAME_PREFIX = "hugegraph_llm_" diff --git a/hugegraph-llm/src/tests/indices/test_faiss_vector_index.py b/hugegraph-llm/src/tests/indices/test_faiss_vector_index.py index 0ea3ed316..fd1eb2a15 100644 --- a/hugegraph-llm/src/tests/indices/test_faiss_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_faiss_vector_index.py @@ -18,8 +18,9 @@ import unittest from pprint import pprint -from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding + from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex +from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding class TestVectorIndex(unittest.TestCase): diff --git a/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py index b339111a4..27c8f2995 100644 --- a/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py @@ -19,8 +19,8 @@ import unittest from pprint import pprint -from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding from hugegraph_llm.indices.vector_index.milvus_vector_store import MilvusVectorIndex +from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding test_name = "test" diff --git a/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py index 53288d662..0de26528d 100644 --- a/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py @@ -18,8 +18,9 @@ import unittest from pprint import pprint -from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding + from hugegraph_llm.indices.vector_index.qdrant_vector_store import QdrantVectorIndex +from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding class TestQdrantVectorIndex(unittest.TestCase): From 1a14a4b3c1f996092caf3380c1468a61c4d474b3 Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Thu, 22 May 2025 02:07:40 +0800 Subject: [PATCH 11/71] feat(llm): vector db finished --- .../src/hugegraph_llm/config/__init__.py | 2 +- .../hugegraph_llm/config/hugegraph_config.py | 1 + .../src/hugegraph_llm/config/index_config.py | 2 + .../src/hugegraph_llm/config/llm_config.py | 100 +++++--- .../demo/rag_demo/configs_block.py | 237 +++++++++++++++++- .../hugegraph_llm/demo/rag_demo/rag_block.py | 69 ++++- .../demo/rag_demo/text2gremlin_block.py | 35 ++- .../demo/rag_demo/vector_graph_block.py | 87 ++++++- .../indices/vector_index/base.py | 43 ++-- .../vector_index/faiss_vector_store.py | 59 ++++- .../vector_index/milvus_vector_store.py | 103 +++++++- .../vector_index/qdrant_vector_store.py | 84 ++++++- .../hugegraph_llm/middleware/middleware.py | 5 + .../hugegraph_llm/models/embeddings/base.py | 14 ++ .../models/embeddings/init_embedding.py | 30 +++ .../models/embeddings/litellm.py | 17 +- .../hugegraph_llm/models/embeddings/ollama.py | 44 +++- .../hugegraph_llm/models/embeddings/openai.py | 14 ++ .../models/embeddings/qianfan.py | 66 +++++ .../hugegraph_llm/operators/graph_rag_task.py | 19 +- .../operators/gremlin_generate_task.py | 16 +- .../index_op/build_gremlin_example_index.py | 21 +- .../index_op/build_semantic_index.py | 56 ++++- .../operators/index_op/build_vector_index.py | 24 +- .../index_op/gremlin_example_index_query.py | 28 ++- .../operators/index_op/semantic_id_query.py | 18 +- .../operators/index_op/vector_index_query.py | 14 +- .../operators/kg_construction_task.py | 17 +- .../operators/llm_op/gremlin_generate.py | 8 +- .../hugegraph_llm/utils/graph_index_utils.py | 30 ++- .../hugegraph_llm/utils/vector_index_utils.py | 50 +++- .../tests/indices/test_milvus_vector_index.py | 2 +- .../tests/indices/test_qdrant_vector_index.py | 2 +- 33 files changed, 1159 insertions(+), 158 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/models/embeddings/qianfan.py diff --git a/hugegraph-llm/src/hugegraph_llm/config/__init__.py b/hugegraph-llm/src/hugegraph_llm/config/__init__.py index 426ceb949..7ff7f2938 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/__init__.py +++ b/hugegraph-llm/src/hugegraph_llm/config/__init__.py @@ -16,7 +16,7 @@ # under the License. -__all__ = ["huge_settings", "admin_settings", "llm_settings", "resource_path"] +__all__ = ["huge_settings", "admin_settings", "llm_settings", "resource_path", "index_settings"] import os diff --git a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py b/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py index 18710bcdf..1937eda10 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py @@ -16,6 +16,7 @@ # under the License. from typing import Optional + from .models import BaseConfig diff --git a/hugegraph-llm/src/hugegraph_llm/config/index_config.py b/hugegraph-llm/src/hugegraph_llm/config/index_config.py index b55353476..21bc509cd 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/index_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/index_config.py @@ -32,3 +32,5 @@ class IndexConfig(BaseConfig): milvus_port: int = int(os.environ.get("MILVUS_PORT", "19530")) milvus_user: str = os.environ.get("MILVUS_USER", "") milvus_password: str = os.environ.get("MILVUS_PASSWORD", "") + + now_vector_index: str = 'Faiss' diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index 304e5ac44..e3e75a66c 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -32,61 +32,81 @@ class LLMConfig(BaseConfig): embedding_type: Optional[Literal["openai", "litellm", "ollama/local"]] = "openai" reranker_type: Optional[Literal["cohere", "siliconflow"]] = None # 1. OpenAI settings - openai_chat_api_base: Optional[str] = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) - openai_chat_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") - openai_chat_language_model: Optional[str] = "gpt-4.1-mini" - openai_extract_api_base: Optional[str] = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) - openai_extract_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") - openai_extract_language_model: Optional[str] = "gpt-4.1-mini" - openai_text2gql_api_base: Optional[str] = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) - openai_text2gql_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") - openai_text2gql_language_model: Optional[str] = "gpt-4.1-mini" - openai_embedding_api_base: Optional[str] = os.environ.get( - "OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1" - ) - openai_embedding_api_key: Optional[str] = os.environ.get("OPENAI_EMBEDDING_API_KEY") - openai_embedding_model: Optional[str] = "text-embedding-3-small" + openai_chat_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_chat_api_key: str | None = os.environ.get("OPENAI_API_KEY") + openai_chat_language_model: str = "gpt-4.1-mini" + openai_extract_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_extract_api_key: str | None = os.environ.get("OPENAI_API_KEY") + openai_extract_language_model: str = "gpt-4.1-mini" + openai_text2gql_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_text2gql_api_key: str | None = os.environ.get("OPENAI_API_KEY") + openai_text2gql_language_model: str = "gpt-4.1-mini" + openai_embedding_api_base: str = os.environ.get("OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1") + openai_embedding_api_key: str | None = os.environ.get("OPENAI_EMBEDDING_API_KEY") + openai_embedding_model: str = "text-embedding-3-small" + openai_embedding_model_dim: int = 1536 openai_chat_tokens: int = 8192 openai_extract_tokens: int = 256 openai_text2gql_tokens: int = 4096 # 2. Rerank settings - cohere_base_url: Optional[str] = os.environ.get( - "CO_API_URL", "https://api.cohere.com/v1/rerank" - ) + cohere_base_url: str = os.environ.get("CO_API_URL", "https://api.cohere.com/v1/rerank") reranker_api_key: Optional[str] = None reranker_model: Optional[str] = None # 3. Ollama settings - ollama_chat_host: Optional[str] = "127.0.0.1" - ollama_chat_port: Optional[int] = 11434 - ollama_chat_language_model: Optional[str] = None - ollama_extract_host: Optional[str] = "127.0.0.1" - ollama_extract_port: Optional[int] = 11434 - ollama_extract_language_model: Optional[str] = None - ollama_text2gql_host: Optional[str] = "127.0.0.1" - ollama_text2gql_port: Optional[int] = 11434 - ollama_text2gql_language_model: Optional[str] = None - ollama_embedding_host: Optional[str] = "127.0.0.1" - ollama_embedding_port: Optional[int] = 11434 - ollama_embedding_model: Optional[str] = None - # 4. LiteLLM settings + ollama_chat_host: str = "127.0.0.1" + ollama_chat_port: int = 11434 + ollama_chat_language_model: str | None = None + ollama_extract_host: str = "127.0.0.1" + ollama_extract_port: int = 11434 + ollama_extract_language_model: str | None = None + ollama_text2gql_host: str = "127.0.0.1" + ollama_text2gql_port: int = 11434 + ollama_text2gql_language_model: str | None = None + ollama_embedding_host: str = os.getenv("OLLAMA_EMBEDDING_HOST", "127.0.0.1") + ollama_embedding_port: int = int(os.getenv("OLLAMA_EMBEDDING_PORT", 11434)) + ollama_embedding_model: str = os.getenv("OLLAMA_EMBEDDING_MODEL", 'quentinz/bge-large-zh-v1.5') + ollama_embedding_model_dim: Optional[int] = ( + int(os.getenv("OLLAMA_EMBEDDING_MODEL_DIM")) if os.getenv("OLLAMA_EMBEDDING_MODEL_DIM") else None # type:ignore + ) + + # 4. QianFan/WenXin settings + # TODO: update to one token key mode + qianfan_chat_api_key: Optional[str] = None + qianfan_chat_secret_key: Optional[str] = None + qianfan_chat_access_token: Optional[str] = None + qianfan_extract_api_key: Optional[str] = None + qianfan_extract_secret_key: Optional[str] = None + qianfan_extract_access_token: Optional[str] = None + qianfan_text2gql_api_key: Optional[str] = None + qianfan_text2gql_secret_key: Optional[str] = None + qianfan_text2gql_access_token: Optional[str] = None + qianfan_embedding_api_key: Optional[str] = None + qianfan_embedding_secret_key: Optional[str] = None + # 4.1 URL settings + qianfan_url_prefix: str = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop" + qianfan_chat_url: str = qianfan_url_prefix + "/chat/" + qianfan_chat_language_model: str = "ERNIE-Speed-128K" + qianfan_extract_language_model: str = "ERNIE-Speed-128K" + qianfan_text2gql_language_model: str = "ERNIE-Speed-128K" + qianfan_embed_url: str = qianfan_url_prefix + "/embeddings/" + qianfan_embedding_model_dim: int = 384 + + # refer https://cloud.baidu.com/doc/WENXINWORKSHOP/s/alj562vvu to get more details + qianfan_embedding_model: str = "embedding-v1" + # 5. LiteLLM settings litellm_chat_api_key: Optional[str] = None litellm_chat_api_base: Optional[str] = None - litellm_chat_language_model: Optional[str] = "openai/gpt-4.1-mini" + litellm_chat_language_model: str = "openai/gpt-4.1-mini" litellm_chat_tokens: int = 8192 litellm_extract_api_key: Optional[str] = None litellm_extract_api_base: Optional[str] = None - litellm_extract_language_model: Optional[str] = "openai/gpt-4.1-mini" + litellm_extract_language_model: str = "openai/gpt-4.1-mini" litellm_extract_tokens: int = 256 litellm_text2gql_api_key: Optional[str] = None litellm_text2gql_api_base: Optional[str] = None - litellm_text2gql_language_model: Optional[str] = "openai/gpt-4.1-mini" + litellm_text2gql_language_model: str = "openai/gpt-4.1-mini" litellm_text2gql_tokens: int = 4096 litellm_embedding_api_key: Optional[str] = None litellm_embedding_api_base: Optional[str] = None - litellm_embedding_model: Optional[str] = "openai/text-embedding-3-small" + litellm_embedding_model: str = "openai/text-embedding-3-small" + litellm_embedding_model_dim: int = 1536 diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 8c595c30d..e2633229e 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -25,7 +25,7 @@ from dotenv import dotenv_values from requests.auth import HTTPBasicAuth -from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.config import huge_settings, index_settings, llm_settings from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding from hugegraph_llm.models.llms.litellm import LiteLLMClient from hugegraph_llm.utils.log import log @@ -33,8 +33,12 @@ current_llm = "chat" -def test_litellm_embedding(api_key, api_base, model_name) -> int: +def test_litellm_embedding(api_key, api_base, model_name, model_dim) -> int: llm_client = LiteLLMEmbedding( +<<<<<<< HEAD +======= + embedding_dimension=model_dim, +>>>>>>> 38dce0b (feat(llm): vector db finished) api_key=api_key, api_base=api_base, model_name=model_name, @@ -106,29 +110,61 @@ def test_api_connection( return resp.status_code +<<<<<<< HEAD def apply_embedding_config(arg1, arg2, arg3, origin_call=None) -> int: +======= +def config_qianfan_model(arg1, arg2, arg3=None, settings_prefix=None, origin_call=None) -> int: + setattr(llm_settings, f"qianfan_{settings_prefix}_api_key", arg1) + setattr(llm_settings, f"qianfan_{settings_prefix}_secret_key", arg2) + if arg3: + setattr(llm_settings, f"qianfan_{settings_prefix}_language_model", arg3) + params = { + "grant_type": "client_credentials", + "client_id": arg1, + "client_secret": arg2, + } + status_code = test_api_connection( + "https://aip.baidubce.com/oauth/2.0/token", "POST", params=params, origin_call=origin_call + ) + return status_code + + +def apply_embedding_config(arg1, arg2, arg3, arg4, origin_call=None) -> int: +>>>>>>> 38dce0b (feat(llm): vector db finished) status_code = -1 embedding_option = llm_settings.embedding_type + arg4 = int(arg4) if embedding_option == "openai": llm_settings.openai_embedding_api_key = arg1 llm_settings.openai_embedding_api_base = arg2 llm_settings.openai_embedding_model = arg3 + llm_settings.openai_embedding_model_dim = arg4 test_url = llm_settings.openai_embedding_api_base + "/embeddings" headers = {"Authorization": f"Bearer {arg1}"} data = {"model": arg3, "input": "test"} +<<<<<<< HEAD status_code = test_api_connection( test_url, method="POST", headers=headers, body=data, origin_call=origin_call ) +======= + status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) + elif embedding_option == "qianfan_wenxin": + status_code = config_qianfan_model(arg1, arg2, settings_prefix="embedding", origin_call=origin_call) + llm_settings.qianfan_embedding_model = arg3 + llm_settings.qianfan_embedding_model_dim = arg4 +>>>>>>> 38dce0b (feat(llm): vector db finished) elif embedding_option == "ollama/local": llm_settings.ollama_embedding_host = arg1 llm_settings.ollama_embedding_port = int(arg2) llm_settings.ollama_embedding_model = arg3 + llm_settings.ollama_embedding_model_dim = arg4 status_code = test_api_connection(f"http://{arg1}:{arg2}", origin_call=origin_call) elif embedding_option == "litellm": llm_settings.litellm_embedding_api_key = arg1 llm_settings.litellm_embedding_api_base = arg2 llm_settings.litellm_embedding_model = arg3 - status_code = test_litellm_embedding(arg1, arg2, arg3) + llm_settings.litellm_embedding_model_dim = arg4 + status_code = test_litellm_embedding(arg1, arg2, arg3, arg4) llm_settings.update_env() gr.Info("Configured!") return status_code @@ -145,10 +181,10 @@ def apply_reranker_config( if reranker_option == "cohere": llm_settings.reranker_api_key = reranker_api_key llm_settings.reranker_model = reranker_model - llm_settings.cohere_base_url = cohere_base_url + llm_settings.cohere_base_url = cohere_base_url # type:ignore headers = {"Authorization": f"Bearer {reranker_api_key}"} status_code = test_api_connection( - cohere_base_url.rsplit("/", 1)[0] + "/check-api-key", + cohere_base_url.rsplit("/", 1)[0] + "/check-api-key", # type:ignore method="POST", headers=headers, origin_call=origin_call, @@ -221,10 +257,20 @@ def apply_llm_config( "temperature": 0.01, "messages": [{"role": "user", "content": "test"}], } +<<<<<<< HEAD headers = {"Authorization": f"Bearer {api_key_or_host}"} status_code = test_api_connection( test_url, method="POST", headers=headers, body=data, origin_call=origin_call ) +======= + headers = {"Authorization": f"Bearer {arg1}"} + status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) + + elif llm_option == "qianfan_wenxin": + status_code = config_qianfan_model( + arg1, arg2, arg3, settings_prefix=current_llm_config, origin_call=origin_call + ) # pylint: disable=C0301 +>>>>>>> 38dce0b (feat(llm): vector db finished) elif llm_option == "ollama/local": setattr(llm_settings, f"ollama_{current_llm_config}_host", api_key_or_host) @@ -262,6 +308,7 @@ def create_configs_block() -> list: info="IP:PORT (e.g. 127.0.0.1:8080) or full URL (e.g. http://127.0.0.1:8080)", ), gr.Textbox( +<<<<<<< HEAD value=huge_settings.graph_name, label="graph", info="The graph name of HugeGraph-Server instance", @@ -276,6 +323,13 @@ def create_configs_block() -> list: label="pwd", type="password", info="Password for graph server auth", +======= + value=huge_settings.graph_name, label="graph", info="The graph name of HugeGraph-Server instance" + ), + gr.Textbox(value=huge_settings.graph_user, label="user", info="Username for graph server auth"), + gr.Textbox( + value=huge_settings.graph_pwd, label="pwd", type="password", info="Password for graph server auth" +>>>>>>> 38dce0b (feat(llm): vector db finished) ), gr.Textbox( value=huge_settings.graph_space, @@ -294,9 +348,15 @@ def create_configs_block() -> list: "> Tips: The OpenAI option also support openai style api from other providers. " "**Refresh the page** to load the **latest configs** in __UI__." ) +<<<<<<< HEAD with gr.Tab(label="chat"): chat_llm_dropdown = gr.Dropdown( choices=["openai", "litellm", "ollama/local"], +======= + with gr.Tab(label='chat'): + chat_llm_dropdown = gr.Dropdown( + choices=["openai", "litellm", "qianfan_wenxin", "ollama/local"], +>>>>>>> 38dce0b (feat(llm): vector db finished) value=getattr(llm_settings, "chat_llm_type"), label="type", ) @@ -308,6 +368,7 @@ def chat_llm_settings(llm_type): if llm_type == "openai": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD value=getattr(llm_settings, "openai_chat_api_key"), label="api_key", type="password", @@ -339,14 +400,42 @@ def chat_llm_settings(llm_type): value=getattr(llm_settings, "ollama_chat_language_model"), label="model_name", ), +======= + value=getattr(llm_settings, "openai_chat_api_key"), label="api_key", type="password" + ), + gr.Textbox(value=getattr(llm_settings, "openai_chat_api_base"), label="api_base"), + gr.Textbox(value=getattr(llm_settings, "openai_chat_language_model"), label="model_name"), + gr.Textbox(value=getattr(llm_settings, "openai_chat_tokens"), label="max_token"), + ] + elif llm_type == "ollama/local": + llm_config_input = [ + gr.Textbox(value=getattr(llm_settings, "ollama_chat_host"), label="host"), + gr.Textbox(value=str(getattr(llm_settings, "ollama_chat_port")), label="port"), + gr.Textbox(value=getattr(llm_settings, "ollama_chat_language_model"), label="model_name"), + gr.Textbox(value="", visible=False), + ] + elif llm_type == "qianfan_wenxin": + llm_config_input = [ + gr.Textbox( + value=getattr(llm_settings, "qianfan_chat_api_key"), label="api_key", type="password" + ), + gr.Textbox( + value=getattr(llm_settings, "qianfan_chat_secret_key"), label="secret_key", type="password" + ), + gr.Textbox(value=getattr(llm_settings, "qianfan_chat_language_model"), label="model_name"), +>>>>>>> 38dce0b (feat(llm): vector db finished) gr.Textbox(value="", visible=False), ] elif llm_type == "litellm": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD value=getattr(llm_settings, "litellm_chat_api_key"), label="api_key", type="password", +======= + value=getattr(llm_settings, "litellm_chat_api_key"), label="api_key", type="password" +>>>>>>> 38dce0b (feat(llm): vector db finished) ), gr.Textbox( value=getattr(llm_settings, "litellm_chat_api_base"), @@ -358,10 +447,14 @@ def chat_llm_settings(llm_type): label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), +<<<<<<< HEAD gr.Textbox( value=getattr(llm_settings, "litellm_chat_tokens"), label="max_token", ), +======= + gr.Textbox(value=getattr(llm_settings, "litellm_chat_tokens"), label="max_token"), +>>>>>>> 38dce0b (feat(llm): vector db finished) ] else: llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] @@ -379,6 +472,7 @@ def chat_llm_settings(llm_type): apply_llm_config_with_text2gql_op, inputs=llm_config_input ) if not api_text2sql_key: +<<<<<<< HEAD llm_config_button.click( apply_llm_config_with_extract_op, inputs=llm_config_input ) @@ -386,6 +480,13 @@ def chat_llm_settings(llm_type): with gr.Tab(label="mini_tasks"): extract_llm_dropdown = gr.Dropdown( choices=["openai", "litellm", "ollama/local"], +======= + llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) + + with gr.Tab(label='mini_tasks'): + extract_llm_dropdown = gr.Dropdown( + choices=["openai", "litellm", "qianfan_wenxin", "ollama/local"], +>>>>>>> 38dce0b (feat(llm): vector db finished) value=getattr(llm_settings, "extract_llm_type"), label="type", ) @@ -397,6 +498,7 @@ def extract_llm_settings(llm_type): if llm_type == "openai": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD value=getattr(llm_settings, "openai_extract_api_key"), label="api_key", type="password", @@ -428,14 +530,44 @@ def extract_llm_settings(llm_type): value=getattr(llm_settings, "ollama_extract_language_model"), label="model_name", ), +======= + value=getattr(llm_settings, "openai_extract_api_key"), label="api_key", type="password" + ), + gr.Textbox(value=getattr(llm_settings, "openai_extract_api_base"), label="api_base"), + gr.Textbox(value=getattr(llm_settings, "openai_extract_language_model"), label="model_name"), + gr.Textbox(value=getattr(llm_settings, "openai_extract_tokens"), label="max_token"), + ] + elif llm_type == "ollama/local": + llm_config_input = [ + gr.Textbox(value=getattr(llm_settings, "ollama_extract_host"), label="host"), + gr.Textbox(value=str(getattr(llm_settings, "ollama_extract_port")), label="port"), + gr.Textbox(value=getattr(llm_settings, "ollama_extract_language_model"), label="model_name"), + gr.Textbox(value="", visible=False), + ] + elif llm_type == "qianfan_wenxin": + llm_config_input = [ + gr.Textbox( + value=getattr(llm_settings, "qianfan_extract_api_key"), label="api_key", type="password" + ), + gr.Textbox( + value=getattr(llm_settings, "qianfan_extract_secret_key"), + label="secret_key", + type="password", + ), + gr.Textbox(value=getattr(llm_settings, "qianfan_extract_language_model"), label="model_name"), +>>>>>>> 38dce0b (feat(llm): vector db finished) gr.Textbox(value="", visible=False), ] elif llm_type == "litellm": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD value=getattr(llm_settings, "litellm_extract_api_key"), label="api_key", type="password", +======= + value=getattr(llm_settings, "litellm_extract_api_key"), label="api_key", type="password" +>>>>>>> 38dce0b (feat(llm): vector db finished) ), gr.Textbox( value=getattr(llm_settings, "litellm_extract_api_base"), @@ -447,19 +579,29 @@ def extract_llm_settings(llm_type): label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), +<<<<<<< HEAD gr.Textbox( value=getattr(llm_settings, "litellm_extract_tokens"), label="max_token", ), +======= + gr.Textbox(value=getattr(llm_settings, "litellm_extract_tokens"), label="max_token"), +>>>>>>> 38dce0b (feat(llm): vector db finished) ] else: llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) +<<<<<<< HEAD with gr.Tab(label="text2gql"): text2gql_llm_dropdown = gr.Dropdown( choices=["openai", "litellm", "ollama/local"], +======= + with gr.Tab(label='text2gql'): + text2gql_llm_dropdown = gr.Dropdown( + choices=["openai", "litellm", "qianfan_wenxin", "ollama/local"], +>>>>>>> 38dce0b (feat(llm): vector db finished) value=getattr(llm_settings, "text2gql_llm_type"), label="type", ) @@ -471,6 +613,7 @@ def text2gql_llm_settings(llm_type): if llm_type == "openai": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD value=getattr(llm_settings, "openai_text2gql_api_key"), label="api_key", type="password", @@ -502,14 +645,44 @@ def text2gql_llm_settings(llm_type): value=getattr(llm_settings, "ollama_text2gql_language_model"), label="model_name", ), +======= + value=getattr(llm_settings, "openai_text2gql_api_key"), label="api_key", type="password" + ), + gr.Textbox(value=getattr(llm_settings, "openai_text2gql_api_base"), label="api_base"), + gr.Textbox(value=getattr(llm_settings, "openai_text2gql_language_model"), label="model_name"), + gr.Textbox(value=getattr(llm_settings, "openai_text2gql_tokens"), label="max_token"), + ] + elif llm_type == "ollama/local": + llm_config_input = [ + gr.Textbox(value=getattr(llm_settings, "ollama_text2gql_host"), label="host"), + gr.Textbox(value=str(getattr(llm_settings, "ollama_text2gql_port")), label="port"), + gr.Textbox(value=getattr(llm_settings, "ollama_text2gql_language_model"), label="model_name"), + gr.Textbox(value="", visible=False), + ] + elif llm_type == "qianfan_wenxin": + llm_config_input = [ + gr.Textbox( + value=getattr(llm_settings, "qianfan_text2gql_api_key"), label="api_key", type="password" + ), + gr.Textbox( + value=getattr(llm_settings, "qianfan_text2gql_secret_key"), + label="secret_key", + type="password", + ), + gr.Textbox(value=getattr(llm_settings, "qianfan_text2gql_language_model"), label="model_name"), +>>>>>>> 38dce0b (feat(llm): vector db finished) gr.Textbox(value="", visible=False), ] elif llm_type == "litellm": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD value=getattr(llm_settings, "litellm_text2gql_api_key"), label="api_key", type="password", +======= + value=getattr(llm_settings, "litellm_text2gql_api_key"), label="api_key", type="password" +>>>>>>> 38dce0b (feat(llm): vector db finished) ), gr.Textbox( value=getattr(llm_settings, "litellm_text2gql_api_base"), @@ -521,10 +694,14 @@ def text2gql_llm_settings(llm_type): label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), +<<<<<<< HEAD gr.Textbox( value=getattr(llm_settings, "litellm_text2gql_tokens"), label="max_token", ), +======= + gr.Textbox(value=getattr(llm_settings, "litellm_text2gql_tokens"), label="max_token"), +>>>>>>> 38dce0b (feat(llm): vector db finished) ] else: llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] @@ -533,7 +710,11 @@ def text2gql_llm_settings(llm_type): with gr.Accordion("3. Set up the Embedding.", open=False): embedding_dropdown = gr.Dropdown( +<<<<<<< HEAD choices=["openai", "litellm", "ollama/local"], +======= + choices=["openai", "litellm", "qianfan_wenxin", "ollama/local"], +>>>>>>> 38dce0b (feat(llm): vector db finished) value=llm_settings.embedding_type, label="Embedding", ) @@ -544,6 +725,7 @@ def embedding_settings(embedding_type): if embedding_type == "openai": with gr.Row(): embedding_config_input = [ +<<<<<<< HEAD gr.Textbox( value=llm_settings.openai_embedding_api_key, label="api_key", @@ -557,24 +739,49 @@ def embedding_settings(embedding_type): value=llm_settings.openai_embedding_model, label="model_name", ), +======= + gr.Textbox(value=llm_settings.openai_embedding_api_key, label="api_key", type="password"), + gr.Textbox(value=llm_settings.openai_embedding_api_base, label="api_base"), + gr.Textbox(value=llm_settings.openai_embedding_model, label="model_name"), + gr.Textbox(value=str(llm_settings.openai_embedding_model_dim), label="model_dim"), +>>>>>>> 38dce0b (feat(llm): vector db finished) ] elif embedding_type == "ollama/local": with gr.Row(): embedding_config_input = [ gr.Textbox(value=llm_settings.ollama_embedding_host, label="host"), gr.Textbox(value=str(llm_settings.ollama_embedding_port), label="port"), +<<<<<<< HEAD gr.Textbox( value=llm_settings.ollama_embedding_model, label="model_name", ), +======= + gr.Textbox(value=llm_settings.ollama_embedding_model, label="model_name"), + gr.Textbox(value=str(llm_settings.ollama_embedding_model_dim), label="model_dim"), + ] + elif embedding_type == "qianfan_wenxin": + with gr.Row(): + embedding_config_input = [ + gr.Textbox(value=llm_settings.qianfan_embedding_api_key, label="api_key", type="password"), + gr.Textbox( + value=llm_settings.qianfan_embedding_secret_key, label="secret_key", type="password" + ), + gr.Textbox(value=llm_settings.qianfan_embedding_model, label="model_name"), + gr.Textbox(value=str(llm_settings.qianfan_embedding_model_dim), label="model_dim"), +>>>>>>> 38dce0b (feat(llm): vector db finished) ] elif embedding_type == "litellm": with gr.Row(): embedding_config_input = [ gr.Textbox( +<<<<<<< HEAD value=getattr(llm_settings, "litellm_embedding_api_key"), label="api_key", type="password", +======= + value=getattr(llm_settings, "litellm_embedding_api_key"), label="api_key", type="password" +>>>>>>> 38dce0b (feat(llm): vector db finished) ), gr.Textbox( value=getattr(llm_settings, "litellm_embedding_api_base"), @@ -586,12 +793,19 @@ def embedding_settings(embedding_type): label="model_name", info="Please refer to https://docs.litellm.ai/docs/embedding/supported_embedding", ), +<<<<<<< HEAD +======= + gr.Textbox( + value=getattr(llm_settings, "litellm_embedding_model_dim"), label="model_dim", type="text" + ), +>>>>>>> 38dce0b (feat(llm): vector db finished) ] else: embedding_config_input = [ gr.Textbox(value="", visible=False), gr.Textbox(value="", visible=False), gr.Textbox(value="", visible=False), + gr.Textbox(value="", visible=False), ] embedding_config_button = gr.Button("Apply Configuration") @@ -653,6 +867,19 @@ def reranker_settings(reranker_type): inputs=reranker_config_input, # pylint: disable=no-member ) +<<<<<<< HEAD +======= + with gr.Accordion("5. Set up the vector database.", open=False): + engine_selector = gr.Dropdown( + choices=["Faiss", "Milvus", "Qdrant"], + value=lambda: index_settings.now_vector_index, + label="Select vector database.", + ) + engine_selector.select( + fn=lambda engine: setattr(index_settings, "now_vector_index", engine), + inputs=[engine_selector], + ) +>>>>>>> 38dce0b (feat(llm): vector db finished) # The reason for returning this partial value is the functional need to refresh the ui return graph_config_input diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 5ff3df931..f61b57123 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -18,14 +18,20 @@ # pylint: disable=E1101 import os -from typing import AsyncGenerator, Tuple, Literal, Optional +from typing import AsyncGenerator, Literal, Optional, Tuple import gradio as gr from hugegraph_llm.flows.scheduler import SchedulerSingleton import pandas as pd from gradio.utils import NamedString +<<<<<<< HEAD from hugegraph_llm.config import resource_path, prompt, llm_settings +======= +from hugegraph_llm.config import huge_settings, index_settings, llm_settings, prompt, resource_path +from hugegraph_llm.operators.graph_rag_task import RAGPipeline +from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize +>>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.utils.decorators import with_task_id from hugegraph_llm.utils.log import log @@ -71,7 +77,32 @@ def rag_answer( gr.Warning("Please select at least one generate mode.") return "", "", "", "" +<<<<<<< HEAD scheduler = SchedulerSingleton.get_instance() +======= + rag = RAGPipeline() + if vector_search: + rag.query_vector_index(vector_index_str=index_settings.now_vector_index) + if graph_search: + rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid( + vector_index_str=index_settings.now_vector_index, + vector_dis_threshold=vector_dis_threshold, + topk_per_keyword=topk_per_keyword, + ).import_schema(huge_settings.graph_name).query_graphdb( + num_gremlin_generate_example=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, + max_graph_items=max_graph_items, + ) + # TODO: add more user-defined search strategies + rag.merge_dedup_rerank( + graph_ratio=graph_ratio, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + topk_return_results=topk_return_results, + ) + rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) + +>>>>>>> 38dce0b (feat(llm): vector db finished) try: # Select workflow by mode to avoid fetching the wrong pipeline from the pool if graph_vector_answer or (graph_only_answer and vector_only_answer): @@ -193,6 +224,26 @@ async def rag_answer_streaming( yield "", "", "", "" return +<<<<<<< HEAD +======= + rag = RAGPipeline() + if vector_search: + rag.query_vector_index(vector_index_str=index_settings.now_vector_index) + if graph_search: + rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid( + vector_index_str=index_settings.now_vector_index + ).import_schema(huge_settings.graph_name).query_graphdb( + num_gremlin_generate_example=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, + ) + rag.merge_dedup_rerank( + graph_ratio, + rerank_method, + near_neighbor_first, + ) + # rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) + +>>>>>>> 38dce0b (feat(llm): vector db finished) try: # Select the specific streaming workflow scheduler = SchedulerSingleton.get_instance() @@ -332,9 +383,7 @@ def toggle_slider(enable): 0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False ) - graph_vector_radio.change( - toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio - ) # pylint: disable=no-member + graph_vector_radio.change(toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio) # pylint: disable=no-member near_neighbor_first = gr.Checkbox( value=False, label="Near neighbor first(Optional)", @@ -397,15 +446,15 @@ def read_file_to_excel(file: NamedString, line_count: Optional[int] = None): df = pd.read_excel(file.name, nrows=line_count) if file else pd.DataFrame() elif file.name.endswith(".csv"): df = pd.read_csv(file.name, nrows=line_count) if file else pd.DataFrame() - df.to_excel(questions_path, index=False) - if df.empty: + df.to_excel(questions_path, index=False) # type:ignore + if df.empty: # type:ignore df = pd.DataFrame([[""] * len(tests_df_headers)], columns=tests_df_headers) else: - df.columns = tests_df_headers + df.columns = tests_df_headers # type:ignore # truncate the dataframe if it's too long - if len(df) > 40: - return df.head(40), 40 - return df, len(df) + if len(df) > 40: # type:ignore + return df.head(40), 40 # type:ignore + return df, len(df) # type:ignore def change_showing_excel(line_count): if os.path.exists(answers_path): diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 6600d7c41..11103d7c4 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -17,14 +17,18 @@ import json import os +<<<<<<< HEAD from datetime import datetime from dataclasses import dataclass from typing import Any, Tuple, Dict, Literal, Optional, List +======= +from typing import Any, Dict, Literal, Tuple, Union +>>>>>>> 38dce0b (feat(llm): vector db finished) import gradio as gr import pandas as pd -from hugegraph_llm.config import prompt, resource_path, huge_settings +from hugegraph_llm.config import huge_settings, index_settings, prompt, resource_path from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.models.llms.init_llm import LLMs from hugegraph_llm.operators.graph_rag_task import RAGPipeline @@ -33,6 +37,7 @@ from hugegraph_llm.utils.embedding_utils import get_index_folder_name from hugegraph_llm.utils.hugegraph_utils import run_gremlin_query from hugegraph_llm.utils.log import log +<<<<<<< HEAD from hugegraph_llm.flows.scheduler import SchedulerSingleton @@ -71,6 +76,9 @@ def success_result( template_exec_result=template_exec, raw_exec_result=raw_exec, ) +======= +from hugegraph_llm.utils.vector_index_utils import get_vector_index_class +>>>>>>> 38dce0b (feat(llm): vector db finished) def store_schema(schema, question, gremlin_prompt): @@ -86,10 +94,15 @@ def store_schema(schema, question, gremlin_prompt): def build_example_vector_index(temp_file) -> dict: +<<<<<<< HEAD folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) index_path = os.path.join(resource_path, folder_name, "gremlin_examples") if not os.path.exists(index_path): os.makedirs(index_path) +======= + vector_index = get_vector_index_class(index_settings.now_vector_index) + assert vector_index, 'vector db name is error' +>>>>>>> 38dce0b (feat(llm): vector db finished) if temp_file is None: full_path = os.path.join(resource_path, "demo", "text2gremlin.csv") else: @@ -120,11 +133,20 @@ def build_example_vector_index(temp_file) -> dict: llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding(), ) - return builder.example_index_build(examples).run() + return builder.example_index_build(examples, vector_index=vector_index).run() +<<<<<<< HEAD def _process_schema(schema, generator, sm): """Process and validate schema input""" +======= +def gremlin_generate( + inp, example_num, schema, gremlin_prompt +) -> Union[tuple[str, str], tuple[str, Any, Any, Any, Any]]: + vector_index = get_vector_index_class(index_settings.now_vector_index) + generator = GremlinGenerator(llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding()) + sm = SchemaManager(graph_name=schema) +>>>>>>> 38dce0b (feat(llm): vector db finished) short_schema = False if not schema: return None, short_schema @@ -200,7 +222,7 @@ def gremlin_generate( output_types = _configure_output_types(requested_outputs) context = ( - generator.example_index_query(example_num) + generator.example_index_query(example_num, vector_index) .gremlin_generate_synthesize(updated_schema, gremlin_prompt) .run(query=inp) ) @@ -220,7 +242,7 @@ def gremlin_generate( def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: - mini_schema = {} + mini_schema = {} # type: ignore # Add necessary vertexlabels items (3) if "vertexlabels" in schema: @@ -299,6 +321,7 @@ def create_text2gremlin_block() -> Tuple: out = gr.Textbox(label="Result Message") with gr.Row(): btn = gr.Button("Build Example Vector Index", variant="primary") + btn.click(build_example_vector_index, inputs=[file], outputs=[out]) # pylint: disable=no-member gr.Markdown("## Nature Language To Gremlin") @@ -364,6 +387,10 @@ def graph_rag_recall( store_schema(prompt.text2gql_graph_schema, query, gremlin_prompt) rag = RAGPipeline() rag.extract_keywords().keywords_to_vid( +<<<<<<< HEAD +======= + vector_index=index_settings.now_vector_index, +>>>>>>> 38dce0b (feat(llm): vector db finished) vector_dis_threshold=vector_dis_threshold, topk_per_keyword=topk_per_keyword, ) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 56b5de4b3..8b081cf9a 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -23,17 +23,21 @@ import gradio as gr +<<<<<<< HEAD from hugegraph_llm.config import huge_settings from hugegraph_llm.config import prompt from hugegraph_llm.config import resource_path from hugegraph_llm.flows.scheduler import SchedulerSingleton +======= +from hugegraph_llm.config import huge_settings, prompt +>>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.utils.graph_index_utils import ( - get_graph_index_info, - clean_all_graph_index, clean_all_graph_data, - update_vid_embedding, + clean_all_graph_index, extract_graph, + get_graph_index_info, import_graph_data, +<<<<<<< HEAD build_schema, ) from hugegraph_llm.utils.hugegraph_utils import check_graph_db_connection @@ -43,6 +47,13 @@ build_vector_index, get_vector_index_info, ) +======= + update_vid_embedding, +) +from hugegraph_llm.utils.hugegraph_utils import check_graph_db_connection +from hugegraph_llm.utils.log import log +from hugegraph_llm.utils.vector_index_utils import build_vector_index, clean_vector_index, get_vector_index_info +>>>>>>> 38dce0b (feat(llm): vector db finished) def store_prompt(doc, schema, example_prompt): @@ -226,6 +237,7 @@ def create_vector_graph_block(): # pylint: disable=no-member # pylint: disable=C0301 # pylint: disable=unexpected-keyword-arg +<<<<<<< HEAD with gr.Blocks() as demo: gr.Markdown( """## Build Vector/Graph Index & Extract Knowledge Graph @@ -240,10 +252,47 @@ def create_vector_graph_block(): - Graph Extract Prompt Header: The user-defined prompt of graph extracting - If already exist the graph data, you should click "**Rebuild vid Index**" to update the index """ +======= + gr.Markdown( + """## Build Vector/Graph Index & Extract Knowledge Graph +- Docs: + - text: Build rag index from plain text + - file: Upload file(s) which should be TXT or .docx (Multiple files can be selected together) +- [Schema](https://hugegraph.apache.org/docs/clients/restful-api/schema/): (Accept **2 types**) + - User-defined Schema (JSON format, follow the [template](https://github.com/apache/incubator-hugegraph-ai/blob/aff3bbe25fa91c3414947a196131be812c20ef11/hugegraph-llm/src/hugegraph_llm/config/config_data.py#L125) + to modify it) + - Specify the name of the HugeGraph graph instance, it will automatically get the schema from it (like + **"hugegraph"**) +- Graph Extract Prompt Header: The user-defined prompt of graph extracting +- If already exist the graph data, you should click "**Rebuild vid Index**" to update the index +""" + ) + + with gr.Row(): + with gr.Column(): + with gr.Tab("text") as tab_upload_text: + input_text = gr.Textbox( + value=prompt.doc_input_text, label="Input Doc(s)", lines=20, show_copy_button=True + ) + with gr.Tab("file") as tab_upload_file: + input_file = gr.File( + value=None, + label="Docs (multi-files can be selected together)", + file_count="multiple", + ) + input_schema = gr.Code(value=prompt.graph_schema, label="Graph Schema", language="json", lines=15, max_lines=29) + info_extract_template = gr.Code( + value=prompt.extract_graph_prompt, + label="Graph Extract Prompt Header", + language="markdown", + lines=15, + max_lines=29, +>>>>>>> 38dce0b (feat(llm): vector db finished) ) with gr.Row(): with gr.Column(): +<<<<<<< HEAD with gr.Tab("text") as tab_upload_text: input_text = gr.Textbox( value=prompt.doc_input_text, @@ -283,12 +332,26 @@ def create_vector_graph_block(): vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") +======= + vector_index_btn0 = gr.Button("Get Vector Index Info", size="sm") + graph_index_btn0 = gr.Button("Get Graph Index Info", size="sm") + with gr.Accordion("Clear RAG Data", open=False): + with gr.Column(): + vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") + graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") + graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") + vector_import_bt = gr.Button("Import into Vector", variant="primary") + graph_extract_bt = gr.Button("Extract Graph Data (1)", variant="primary") + graph_loading_bt = gr.Button("Load into GraphDB (2)", interactive=True) + graph_index_rebuild_bt = gr.Button("Update Vid Embedding") +>>>>>>> 38dce0b (feat(llm): vector db finished) vector_import_bt = gr.Button("Import into Vector", variant="primary") graph_extract_bt = gr.Button("Extract Graph Data (1)", variant="primary") graph_loading_bt = gr.Button("Load into GraphDB (2)", interactive=True) graph_index_rebuild_bt = gr.Button("Update Vid Embedding") +<<<<<<< HEAD gr.Markdown("---") with gr.Accordion("Graph Schema Generator", open=False): gr.Markdown( @@ -343,6 +406,22 @@ def create_vector_graph_block(): store_prompt, inputs=[input_text, input_schema, info_extract_template], ) +======= + # origin_out = gr.Textbox(visible=False) + graph_extract_bt.click( + extract_graph, inputs=[input_file, input_text, input_schema, info_extract_template], outputs=[out] + ).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) + + graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( + update_vid_embedding + ).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) +>>>>>>> 38dce0b (feat(llm): vector db finished) # origin_out = gr.Textbox(visible=False) graph_extract_bt.click( @@ -414,7 +493,7 @@ async def timely_update_vid_embedding(interval_seconds: int = 3600): "pwd": huge_settings.graph_pwd, "graph_space": huge_settings.graph_space, } - if check_graph_db_connection(**config): + if check_graph_db_connection(**config): # type:ignore await asyncio.to_thread(update_vid_embedding) log.info("update_vid_embedding executed successfully") else: diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py index 53743b2eb..feda7a24c 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py @@ -17,7 +17,7 @@ from abc import ABC, abstractmethod -from typing import Any, List, Set, Union +from typing import Any, Dict, List, Set, Union class VectorStoreBase(ABC): @@ -37,6 +37,12 @@ def add(self, vectors: List[List[float]], props: List[Any]): props (List[Any]): List of associated metadata or properties for each vector. """ + @abstractmethod + def get_all_properties(self) -> list[str]: + """ + #TODO: finish comment + """ + @abstractmethod def remove(self, props: Union[Set[Any], List[Any]]) -> int: """ @@ -64,33 +70,36 @@ def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0 """ @abstractmethod - def to_index_file(self, name: str): + def save_index_by_name(self, *name: str): + """ + #TODO: finish comment """ - Persist the vector store (index and metadata) to the specified directory. - Args: - name (str): Path to the directory where the index and properties will be saved. + @abstractmethod + def get_vector_index_info( + self, + ) -> Dict: + """ + #TODO: finish comment """ @staticmethod @abstractmethod - def from_name(name: str) -> "VectorStoreBase": + def from_name(embed_dim: int, *name: str) -> "VectorStoreBase": """ - Load a vector store from the specified directory. - - Args: - name (str): Path to the directory containing the index and properties. - - Returns: - VectorStore: An instance of the vector store. + #TODO: finish comment """ @staticmethod @abstractmethod - def clean(name: str): + def exist(*name: str) -> bool: + """ + #TODO: finish comment """ - Delete the persisted index and properties from the specified directory. - Args: - name (str): Path to the directory to clean. + @staticmethod + @abstractmethod + def clean(*name: str) -> bool: + """ + #TODO: finish comment """ diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py index 934f4b719..ef731db39 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py @@ -18,11 +18,12 @@ import os import pickle as pkl from copy import deepcopy -from typing import Any, List, Set, Union +from typing import Any, Dict, List, Set, Union import faiss import numpy as np +from hugegraph_llm.config import resource_path from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.utils.log import log @@ -35,6 +36,7 @@ def __init__(self, embed_dim: int = 1024): self.index = faiss.IndexFlatL2(embed_dim) self.properties: list[Any] = [] +<<<<<<< HEAD def to_index_file(self, dir_path: str, filename_prefix: str = None): @@ -48,6 +50,12 @@ def to_index_file(self, dir_path: str, filename_prefix: str = None): ) index_file = os.path.join(dir_path, index_name) properties_file = os.path.join(dir_path, property_name) +======= + def save_index_by_name(self, *name: str): + os.makedirs(os.path.join(resource_path, *name), exist_ok=True) + index_file = os.path.join(resource_path, *name, INDEX_FILE_NAME) + properties_file = os.path.join(resource_path, *name, PROPERTIES_FILE_NAME) +>>>>>>> 38dce0b (feat(llm): vector db finished) faiss.write_index(self.index, index_file) with open(properties_file, "wb") as f: pkl.dump(self.properties, f) @@ -55,7 +63,6 @@ def to_index_file(self, dir_path: str, filename_prefix: str = None): def add(self, vectors: List[List[float]], props: List[Any]): if len(vectors) == 0: return - if self.index.ntotal == 0 and len(vectors[0]) != self.index.d: self.index = faiss.IndexFlatL2(len(vectors[0])) self.index.add(np.array(vectors)) @@ -98,7 +105,23 @@ def search( ) return results + def get_all_properties(self) -> list[Any]: + return self.properties + + def get_vector_index_info( + self, + ) -> Dict: + return { + "embed_dim": self.index.d, + "vector_info": { + "chunk_vector_num": self.index.ntotal, + "graph_vid_vector_num": self.index.ntotal, + "graph_properties_vector_num": len(self.properties), + }, + } + @staticmethod +<<<<<<< HEAD def from_index_file( dir_path: str, filename_prefix: str | None = None, record_miss: bool = True ) -> "FaissVectorIndex": @@ -170,20 +193,38 @@ def clean(dir_path: str, filename_prefix: str = None): except OSError as e: log.error("Error removing file %s: %s", file, e) +======= + def clean(*name: str): + index_file = os.path.join(resource_path, *name, INDEX_FILE_NAME) + properties_file = os.path.join(resource_path, *name, PROPERTIES_FILE_NAME) + if os.path.exists(index_file): + os.remove(index_file) + if os.path.exists(properties_file): + os.remove(properties_file) +>>>>>>> 38dce0b (feat(llm): vector db finished) @staticmethod - def from_name(name: str) -> "FaissVectorIndex": - index_file = os.path.join(name, INDEX_FILE_NAME) - properties_file = os.path.join(name, PROPERTIES_FILE_NAME) + def from_name(embed_dim: int, *name: str) -> "FaissVectorIndex": + index_file = os.path.join(resource_path, *name, INDEX_FILE_NAME) + properties_file = os.path.join(resource_path, *name, PROPERTIES_FILE_NAME) if not os.path.exists(index_file) or not os.path.exists(properties_file): log.warning("No index file found, create a new one.") - return FaissVectorIndex() + return FaissVectorIndex(embed_dim) faiss_index = faiss.read_index(index_file) - embed_dim = faiss_index.d with open(properties_file, "rb") as f: properties = pkl.load(f) vector_index = FaissVectorIndex(embed_dim) - vector_index.index = faiss_index - vector_index.properties = properties + if faiss_index.d == vector_index.index.d: + # when dim same, use old + vector_index.index = faiss_index + vector_index.properties = properties + else: + log.warning("dim is different, create a new one.") return vector_index + + @staticmethod + def exist(*name: str) -> bool: + index_file = os.path.join(resource_path, *name, INDEX_FILE_NAME) + properties_file = os.path.join(resource_path, *name, PROPERTIES_FILE_NAME) + return os.path.exists(index_file) and os.path.exists(properties_file) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py index 505209e16..acdbdf511 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -52,6 +52,22 @@ def __init__( if not utility.has_collection(self.name): self._create_collection() + else: + # dim is different, recreate + existing_collection = Collection(self.name) + existing_schema = existing_collection.schema + for field in existing_schema.fields: + if field.name == "embedding" and field.params.get("dim"): + existing_dim = int(field.params["dim"]) + if existing_dim != self.embed_dim: + log.debug( + "Milvus collection '%s' dimension mismatch: %d != %d. Recreating.", + self.name, + existing_dim, + self.embed_dim, + ) + utility.drop_collection(self.name) + break self.collection = Collection(self.name) @@ -76,15 +92,22 @@ def _create_collection(self): } collection.create_index(field_name="embedding", index_params=index_params) - def to_index_file(self, name: str): + def save_index_by_name(self, *name: str): self.collection.flush() - def _deserialize_property(self, prop_str): - """Deserialize property from JSON string.""" + def _deserialize_property(self, prop) -> str: + """If input is a string, return as-is. If dict or list, convert to JSON string.""" + if isinstance(prop, str): + return prop + return json.dumps(prop) + + def _serialize_property(self, prop: str): + """If input is a JSON string, parse it. Otherwise, return as-is.""" try: - return json.loads(prop_str) + return json.loads(prop) except (json.JSONDecodeError, TypeError): - return prop_str + # a simple string + return prop def add(self, vectors: List[List[float]], props: List[Any]): if len(vectors) == 0: @@ -99,7 +122,7 @@ def add(self, vectors: List[List[float]], props: List[Any]): entities.append( { "embedding": vector, - "property": prop, + "property": self._deserialize_property(prop), "original_id": idx, } ) @@ -114,7 +137,7 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: self.collection.load() remove_num = 0 for prop in props: - expr = f'property == "{prop}"' + expr = f'property == "{self._deserialize_property(prop)}"' res = self.collection.delete(expr) if hasattr(res, "delete_count"): remove_num += res.delete_count @@ -144,7 +167,7 @@ def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0 for hit in hits: if hit.distance < dis_threshold: prop_str = hit.entity.get("property") - prop = self._deserialize_property(prop_str) + prop = self._serialize_property(prop_str) ret.append(prop) log.debug("[✓] Add valid distance %s to results.", hit.distance) else: @@ -159,24 +182,78 @@ def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0 finally: self.collection.release() + def get_all_properties(self) -> list[str]: + if self.collection.num_entities == 0: + return [] + + self.collection.load() + try: + results = self.collection.query( + expr='property != ""', + output_fields=["property"], + ) + + return [self._deserialize_property(item["property"]) for item in results] + + finally: + self.collection.release() + + def get_vector_index_info(self) -> dict: + self.collection.load() + try: + embed_dim = None + for field in self.collection.schema.fields: + if field.name == "embedding" and field.dtype == DataType.FLOAT_VECTOR: + embed_dim = int(field.params["dim"]) + break + + if embed_dim is None: + raise ValueError("Could not determine embedding dimension from schema.") + + properties = self.get_all_properties() + return { + "embed_dim": embed_dim, + "vector_info": { + "chunk_vector_num": self.collection.num_entities, + "graph_vid_vector_num": self.collection.num_entities, + "graph_properties_vector_num": len(properties), + }, + } + finally: + self.collection.release() + @staticmethod - def clean(name: str): + def clean(*name: str): + name_str = '_'.join(name) connections.connect( host=index_settings.milvus_host, port=index_settings.milvus_port, user=index_settings.milvus_user, password=index_settings.milvus_password, ) - if utility.has_collection(COLLECTION_NAME_PREFIX + name): - utility.drop_collection(COLLECTION_NAME_PREFIX + name) + if utility.has_collection(COLLECTION_NAME_PREFIX + name_str): + utility.drop_collection(COLLECTION_NAME_PREFIX + name_str) @staticmethod - def from_name(name: str) -> "MilvusVectorIndex": + def from_name(embed_dim: int, *name: str) -> "MilvusVectorIndex": + name_str = '_'.join(name) assert index_settings.milvus_host, "Qdrant host is not configured" return MilvusVectorIndex( - name, + name_str, + host=index_settings.milvus_host, + port=index_settings.milvus_port, + user=index_settings.milvus_user, + password=index_settings.milvus_password, + embed_dim=embed_dim, + ) + + @staticmethod + def exist(*name: str) -> bool: + name_str = '_'.join(name) + connections.connect( host=index_settings.milvus_host, port=index_settings.milvus_port, user=index_settings.milvus_user, password=index_settings.milvus_password, ) + return utility.has_collection(COLLECTION_NAME_PREFIX + name_str) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index 7f392b184..98f2eaf9c 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -from typing import Any, List, Set, Union +from typing import Any, Dict, List, Set, Union from qdrant_client import QdrantClient from qdrant_client.http import models @@ -38,6 +38,18 @@ def __init__(self, name: str, host: str, port: int, api_key=None, embed_dim: int collection_names = [collection.name for collection in collections] if self.name not in collection_names: self._create_collection() + else: + collection_info = self.client.get_collection(self.name) + existing_dim = collection_info.config.params.vectors.size # type: ignore + if existing_dim != self.embed_dim: + log.debug( + "Qdrant collection '%s' dimension mismatch: %d != %d. Recreating.", + self.name, + existing_dim, + self.embed_dim, + ) + self.client.delete_collection(self.name) + self._create_collection() def _create_collection(self): """Create a new collection in Qdrant.""" @@ -47,7 +59,7 @@ def _create_collection(self): ) log.info("Created Qdrant collection '%s'", self.name) - def to_index_file(self, name: str): + def save_index_by_name(self, *name: str): # nothing to do when qdrant pass @@ -123,23 +135,79 @@ def search(self, query_vector: List[float], top_k: int = 5, dis_threshold: float return result_properties + def get_all_properties(self) -> list[str]: + all_properties = [] + offset = None + page_size = 100 + while True: + scroll_result = self.client.scroll( + collection_name=self.name, + offset=offset, + limit=page_size, + with_payload=True, + with_vectors=False, + ) + + points, next_offset = scroll_result + + for point in points: + payload = point.payload + if payload and "property" in payload: + all_properties.append(payload["property"]) + + if next_offset is None or not points: + break + + offset = next_offset + + return all_properties + + def get_vector_index_info(self) -> Dict: + collection_info = self.client.get_collection(self.name) + points_count = collection_info.points_count + embed_dim = collection_info.config.params.vectors.size # type: ignore + + all_properties = self.get_all_properties() + return { + "embed_dim": embed_dim, + "vector_info": { + "chunk_vector_num": points_count, + "graph_vid_vector_num": points_count, + "graph_properties_vector_num": len(all_properties), + }, + } + @staticmethod - def clean(name: str): + def clean(*name: str): + name_str = '_'.join(name) client = QdrantClient( host=index_settings.qdrant_host, port=index_settings.qdrant_port, api_key=index_settings.qdrant_api_key ) collections = client.get_collections().collections collection_names = [collection.name for collection in collections] - name = COLLECTION_NAME_PREFIX + name - if name in collection_names: - client.delete_collection(collection_name=name) + name_str = COLLECTION_NAME_PREFIX + name_str + if name_str in collection_names: + client.delete_collection(collection_name=name_str) @staticmethod - def from_name(name: str) -> "QdrantVectorIndex": + def from_name(embed_dim: int, *name: str) -> "QdrantVectorIndex": assert index_settings.qdrant_host, "Qdrant host is not configured" + name_str = '_'.join(name) return QdrantVectorIndex( - name=COLLECTION_NAME_PREFIX + name, + name=name_str, host=index_settings.qdrant_host, port=index_settings.qdrant_port, + embed_dim=embed_dim, api_key=index_settings.qdrant_api_key, ) + + @staticmethod + def exist(*name: str) -> bool: + name_str = '_'.join(name) + client = QdrantClient( + host=index_settings.qdrant_host, port=index_settings.qdrant_port, api_key=index_settings.qdrant_api_key + ) + collections = client.get_collections().collections + collection_names = [collection.name for collection in collections] + name_str = COLLECTION_NAME_PREFIX + name_str + return name_str in collection_names diff --git a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py index c73242012..7e93d0ec0 100644 --- a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py +++ b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py @@ -46,7 +46,12 @@ async def dispatch(self, request: Request, call_next): "%s - Args: %s, IP: %s, URL: %s", request.method, request.query_params, +<<<<<<< HEAD request.client.host, request.url, +======= + request.client.host, # type: ignore + request.url +>>>>>>> 38dce0b (feat(llm): vector db finished) ) return response diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py index 698b92837..f895314f3 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py @@ -61,6 +61,15 @@ def get_text_embedding(self, text: str) -> List[float]: """Comment""" @abstractmethod +<<<<<<< HEAD +======= + def get_embedding_dim( + self, + ) -> int: + """Comment""" + + @abstractmethod +>>>>>>> 38dce0b (feat(llm): vector db finished) def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch. @@ -81,6 +90,7 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """ @abstractmethod +<<<<<<< HEAD async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch asynchronously. @@ -99,6 +109,10 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] A list of embedding vectors, where each vector is a list of floats. The order of embeddings should match the order of input texts. """ +======= + async def async_get_text_embedding(self, text: str) -> List[float]: + """Comment""" +>>>>>>> 38dce0b (feat(llm): vector db finished) @staticmethod def similarity( diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index 3ad50b3ec..8a96d7182 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -17,6 +17,7 @@ from hugegraph_llm.config import llm_settings +<<<<<<< HEAD from hugegraph_llm.config import LLMConfig from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding @@ -50,6 +51,12 @@ def get_embedding(llm_settings: LLMConfig): ) raise Exception("embedding type is not supported !") +======= +from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding +from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding +from hugegraph_llm.models.embeddings.openai import OpenAIEmbedding +from hugegraph_llm.models.embeddings.qianfan import QianFanEmbedding +>>>>>>> 38dce0b (feat(llm): vector db finished) class Embeddings: @@ -58,22 +65,45 @@ def __init__(self): def get_embedding(self): if self.embedding_type == "openai": + assert llm_settings.openai_embedding_model_dim, 'openai_embedding_model_dim is need' return OpenAIEmbedding( + embedding_dimension=llm_settings.openai_embedding_model_dim, model_name=llm_settings.openai_embedding_model, api_key=llm_settings.openai_embedding_api_key, api_base=llm_settings.openai_embedding_api_base, ) if self.embedding_type == "ollama/local": + assert llm_settings.ollama_embedding_model_dim, 'ollama_embedding_model_dim is need' return OllamaEmbedding( +<<<<<<< HEAD model_name=llm_settings.ollama_embedding_model, host=llm_settings.ollama_embedding_host, port=llm_settings.ollama_embedding_port, ) +======= + embedding_dimension=llm_settings.ollama_embedding_model_dim, + model=llm_settings.ollama_embedding_model, + host=llm_settings.ollama_embedding_host, + port=llm_settings.ollama_embedding_port, + ) + if self.embedding_type == "qianfan_wenxin": + return QianFanEmbedding( + embedding_dimension=llm_settings.litellm_embedding_model_dim, + model_name=llm_settings.qianfan_embedding_model, + api_key=llm_settings.qianfan_embedding_api_key, + secret_key=llm_settings.qianfan_embedding_secret_key, + ) # type: ignore +>>>>>>> 38dce0b (feat(llm): vector db finished) if self.embedding_type == "litellm": return LiteLLMEmbedding( + embedding_dimension=llm_settings.litellm_embedding_model_dim, model_name=llm_settings.litellm_embedding_model, api_key=llm_settings.litellm_embedding_api_key, api_base=llm_settings.litellm_embedding_api_base, +<<<<<<< HEAD ) +======= + ) # type: ignore +>>>>>>> 38dce0b (feat(llm): vector db finished) raise Exception("embedding type is not supported !") diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py index 9d15daa0a..ac0522b99 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py @@ -17,13 +17,17 @@ from typing import List, Optional +<<<<<<< HEAD from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.utils.log import log +======= +from litellm import APIConnectionError, APIError, RateLimitError, aembedding, embedding +>>>>>>> 38dce0b (feat(llm): vector db finished) from tenacity import ( retry, + retry_if_exception_type, stop_after_attempt, wait_exponential, - retry_if_exception_type, ) from litellm import embedding, RateLimitError, APIError, APIConnectionError, aembedding @@ -34,13 +38,24 @@ class LiteLLMEmbedding(BaseEmbedding): def __init__( self, + embedding_dimension, api_key: Optional[str] = None, api_base: Optional[str] = None, model_name: str = "openai/text-embedding-3-small", # Can be any embedding model supported by LiteLLM ) -> None: self.api_key = api_key self.api_base = api_base +<<<<<<< HEAD self.model_name = model_name +======= + self.model = model_name + self.embedding_dimension = embedding_dimension + + def get_embedding_dim( + self, + ) -> int: + return self.embedding_dimension +>>>>>>> 38dce0b (feat(llm): vector db finished) @retry( stop=stop_after_attempt(3), diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py index a4a8bb098..da99d1e65 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py @@ -23,19 +23,54 @@ class OllamaEmbedding(BaseEmbedding): +<<<<<<< HEAD def __init__(self, model_name: str, host: str = "127.0.0.1", port: int = 11434, **kwargs): self.model_name = model_name +======= + def __init__( + self, + model: str = 'quentinz/bge-large-zh-v1.5', + embedding_dimension: int = 1024, + host: str = "127.0.0.1", + port: int = 11434, + **kwargs, + ): + self.model = model +>>>>>>> 38dce0b (feat(llm): vector db finished) self.client = ollama.Client(host=f"http://{host}:{port}", **kwargs) self.async_client = ollama.AsyncClient(host=f"http://{host}:{port}", **kwargs) - self.embedding_dimension = None + self.embedding_dimension = embedding_dimension +<<<<<<< HEAD def get_text_embedding(self, text: str) -> List[float]: """Get embedding for a single text.""" return self.get_texts_embeddings([text])[0] +======= + def get_embedding_dim( + self, + ) -> int: + return self.embedding_dimension + + def get_text_embedding(self, text: str) -> List[float]: + """Comment""" + return list(self.client.embed(model=self.model, input=text)["embeddings"][0]) +>>>>>>> 38dce0b (feat(llm): vector db finished) def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch. +<<<<<<< HEAD +======= + This method efficiently processes multiple texts at once by leveraging + Ollama's batching capabilities, which is more efficient than processing + texts individually. + + Parameters + ---------- + texts : List[str] + A list of text strings to be embedded. + +>>>>>>> 38dce0b (feat(llm): vector db finished) Returns ------- List[List[float]] @@ -52,6 +87,7 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: response = self.client.embed(model=self.model_name, input=texts)["embeddings"] return [list(inner_sequence) for inner_sequence in response] +<<<<<<< HEAD async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch asynchronously. @@ -69,3 +105,9 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] raise AttributeError(error_message) response = await self.async_client.embed(model=self.model_name, input=texts) return [list(inner_sequence) for inner_sequence in response["embeddings"]] +======= + async def async_get_text_embedding(self, text: str) -> List[float]: + """Comment""" + response = await self.async_client.embeddings(model=self.model, prompt=text) + return list(response["embedding"]) +>>>>>>> 38dce0b (feat(llm): vector db finished) diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index d0e15f000..6e30cca71 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -24,6 +24,10 @@ class OpenAIEmbedding: def __init__( self, +<<<<<<< HEAD +======= + embedding_dimension: int = 1536, +>>>>>>> 38dce0b (feat(llm): vector db finished) model_name: str = "text-embedding-3-small", api_key: Optional[str] = None, api_base: Optional[str] = None, @@ -31,7 +35,17 @@ def __init__( api_key = api_key or "" self.client = OpenAI(api_key=api_key, base_url=api_base) self.aclient = AsyncOpenAI(api_key=api_key, base_url=api_base) +<<<<<<< HEAD self.model_name = model_name +======= + self.embedding_model_name = model_name + self.embedding_dimension = embedding_dimension + + def get_embedding_dim( + self, + ) -> int: + return self.embedding_dimension +>>>>>>> 38dce0b (feat(llm): vector db finished) def get_text_embedding(self, text: str) -> List[float]: """Comment""" diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/qianfan.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/qianfan.py new file mode 100644 index 000000000..e5d5463ef --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/qianfan.py @@ -0,0 +1,66 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +from typing import Optional, List + +import qianfan + +from hugegraph_llm.config import llm_settings + +""" +"QianFan" platform can be understood as a unified LLM platform that encompasses the +WenXin large model along with other +common open-source models. + +It enables the invocation and switching between WenXin and these open-source models. +""" + + +class QianFanEmbedding: + def __init__( + self, + embedding_dimension: int, + model_name: str = "embedding-v1", + api_key: Optional[str] = None, + secret_key: Optional[str] = None, + ): + qianfan.get_config().AK = api_key or llm_settings.qianfan_embedding_api_key + qianfan.get_config().SK = secret_key or llm_settings.qianfan_embedding_secret_key + self.embedding_model_name = model_name + self.client = qianfan.Embedding() + self.embedding_dimension = embedding_dimension + + def get_embedding_dim( + self, + ) -> int: + return self.embedding_dimension + + def get_text_embedding(self, text: str) -> List[float]: + """Usage refer: https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlmokk9qn""" + response = self.client.do(model=self.embedding_model_name, texts=[text]) + return response["body"]["data"][0]["embedding"] + + def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: + """Usage refer: https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlmokk9qn""" + response = self.client.do(model=self.embedding_model_name, texts=texts) + return [data["embedding"] for data in response["body"]["data"]] + + async def async_get_text_embedding(self, text: str) -> List[float]: + """Usage refer: https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlmokk9qn""" + response = await self.client.ado(model=self.embedding_model_name, texts=[text]) + return response["body"]["data"][0]["embedding"] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index 330890b5d..80da2b57e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -16,8 +16,10 @@ # under the License. -from typing import Dict, Any, Optional, List, Literal +from typing import Any, Dict, List, Literal, Optional +from hugegraph_llm.config import huge_settings, index_settings, prompt +from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.models.llms.base import BaseLLM @@ -31,8 +33,8 @@ from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize from hugegraph_llm.operators.llm_op.keyword_extract import KeywordExtract -from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm -from hugegraph_llm.config import prompt, huge_settings +from hugegraph_llm.utils.decorators import log_operator_time, log_time, record_rpm +from hugegraph_llm.utils.vector_index_utils import get_vector_index_class class RAGPipeline: @@ -97,6 +99,7 @@ def import_schema(self, graph_name: str): def keywords_to_vid( self, + vector_index_str, by: Literal["query", "keywords"] = "keywords", topk_per_keyword: int = huge_settings.topk_per_keyword, topk_per_query: int = 10, @@ -110,8 +113,10 @@ def keywords_to_vid( :param vector_dis_threshold: Vector distance threshold. :return: Self-instance for chaining. """ + vector_index = get_vector_index_class(vector_index_str=vector_index_str) self._operators.append( SemanticIdQuery( + vector_index=vector_index, embedding=self._embedding, by=by, topk_per_keyword=topk_per_keyword, @@ -156,15 +161,17 @@ def query_graphdb( ) return self - def query_vector_index(self, max_items: int = 3): + def query_vector_index(self, vector_index_str: str, max_items: int = 3): """ Add a vector index query operator to the pipeline. :param max_items: Maximum number of items to retrieve. :return: Self-instance for chaining. """ + vector_index = get_vector_index_class(vector_index_str) self._operators.append( VectorIndexQuery( + vector_index=vector_index, embedding=self._embedding, topk=max_items, ) @@ -244,9 +251,13 @@ def run(self, **kwargs) -> Dict[str, Any]: :return: Final context after all operators have been executed. """ if len(self._operators) == 0: +<<<<<<< HEAD self.extract_keywords().query_graphdb( max_graph_items=kwargs.get("max_graph_items") ).synthesize_answer() +======= + self.extract_keywords().query_graphdb(max_graph_items=kwargs.get('max_graph_items')).synthesize_answer() +>>>>>>> 38dce0b (feat(llm): vector db finished) context = kwargs diff --git a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py index 70f3d27d2..7f8c28fa8 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py @@ -14,8 +14,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from typing import Optional, List +from typing import List, Optional +from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.operators.common_op.check_schema import CheckSchema @@ -24,23 +25,22 @@ from hugegraph_llm.operators.index_op.build_gremlin_example_index import BuildGremlinExampleIndex from hugegraph_llm.operators.index_op.gremlin_example_index_query import GremlinExampleIndexQuery from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize -from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm +from hugegraph_llm.utils.decorators import log_operator_time, log_time, record_rpm class GremlinGenerator: def __init__(self, llm: BaseLLM, embedding: BaseEmbedding): - self.embedding = [] self.llm = llm self.embedding = embedding self.result = None - self.operators = [] + self.operators = [] # type: ignore def clear(self): self.operators = [] return self - def example_index_build(self, examples): - self.operators.append(BuildGremlinExampleIndex(self.embedding, examples)) + def example_index_build(self, examples, vector_index: type[VectorStoreBase]): + self.operators.append(BuildGremlinExampleIndex(self.embedding, examples, vector_index=vector_index)) return self def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): @@ -54,8 +54,8 @@ def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_def raise ValueError("No input data / invalid schema type") return self - def example_index_query(self, num_examples): - self.operators.append(GremlinExampleIndexQuery(self.embedding, num_examples)) + def example_index_query(self, num_examples, vector_index: type[VectorStoreBase]): + self.operators.append(GremlinExampleIndexQuery(vector_index, self.embedding, num_examples)) return self def gremlin_generate_synthesize( diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index fb5c16acb..4f288a9de 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -16,6 +16,7 @@ # under the License. +<<<<<<< HEAD import asyncio import os from typing import Dict, Any, List @@ -28,10 +29,17 @@ get_filename_prefix, get_index_folder_name, ) +======= +from typing import Any, Dict, List + +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +from hugegraph_llm.models.embeddings.base import BaseEmbedding +>>>>>>> 38dce0b (feat(llm): vector db finished) # FIXME: we need keep the logic same with build_semantic_index.py class BuildGremlinExampleIndex: +<<<<<<< HEAD def __init__(self, embedding: BaseEmbedding, examples: List[Dict[str, str]]): self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space @@ -42,6 +50,13 @@ def __init__(self, embedding: BaseEmbedding, examples: List[Dict[str, str]]): self.filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(embedding, "model_name", None) ) +======= + def __init__(self, embedding: BaseEmbedding, examples: List[Dict[str, str]], vector_index: type[VectorStoreBase]): + self.vector_index_name = "gremlin_examples" + self.examples = examples + self.embedding = embedding + self.vector_index = vector_index +>>>>>>> 38dce0b (feat(llm): vector db finished) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: # !: We have assumed that self.example is not empty @@ -50,8 +65,12 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: examples_embedding = asyncio.run(get_embeddings_parallel(self.embedding, queries)) embed_dim = len(examples_embedding[0]) if len(self.examples) > 0: - vector_index = FaissVectorIndex(embed_dim) + vector_index = self.vector_index.from_name(embed_dim, self.vector_index_name) vector_index.add(examples_embedding, self.examples) +<<<<<<< HEAD vector_index.to_index_file(self.index_dir, self.filename_prefix) +======= + vector_index.save_index_by_name(self.vector_index_name) +>>>>>>> 38dce0b (feat(llm): vector db finished) context["embed_dim"] = embed_dim return context 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 ee6837f96..14064af18 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 @@ -17,9 +17,9 @@ import asyncio -import os from typing import Any, Dict +<<<<<<< HEAD from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex from hugegraph_llm.models.embeddings.base import BaseEmbedding @@ -29,10 +29,19 @@ get_filename_prefix, get_index_folder_name, ) +======= +from tqdm import tqdm + +from hugegraph_llm.config import huge_settings +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +from hugegraph_llm.models.embeddings.base import BaseEmbedding +from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager +>>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.utils.log import log class BuildSemanticIndex: +<<<<<<< HEAD def __init__(self, embedding: BaseEmbedding): self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space @@ -42,19 +51,55 @@ def __init__(self, embedding: BaseEmbedding): llm_settings.embedding_type, getattr(embedding, "model_name", None) ) self.vid_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) +======= + 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") +>>>>>>> 38dce0b (feat(llm): vector db finished) self.embedding = embedding self.sm = SchemaManager(huge_settings.graph_name) def _extract_names(self, vertices: list[str]) -> list[str]: return [v.split(":")[1] for v in vertices] +<<<<<<< HEAD +======= + async def _get_embeddings_parallel(self, vids: list[str]) -> list[Any]: + sem = asyncio.Semaphore(10) + batch_size = 1000 + + async def get_embeddings_with_semaphore(vid_list: list[str]) -> Any: + # Executes sync embedding method in a thread pool via loop.run_in_executor, combining async programming + # with multi-threading capabilities. + # This pattern avoids blocking the event loop and prepares for a future fully async pipeline. + async with sem: + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, self.embedding.get_texts_embeddings, vid_list) + + # Split vids into batches of size batch_size + vid_batches = [vids[i : i + batch_size] for i in range(0, len(vids), batch_size)] + + # Create tasks for each batch + tasks = [get_embeddings_with_semaphore(batch) for batch in vid_batches] + + embeddings = [] + with tqdm(total=len(tasks)) as pbar: + for future in asyncio.as_completed(tasks): + batch_embeddings = await future + embeddings.extend(batch_embeddings) # Extend the list with batch results + pbar.update(1) + return embeddings + +>>>>>>> 38dce0b (feat(llm): vector db finished) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: vertexlabels = self.sm.schema.getSchema()["vertexlabels"] all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) - past_vids = self.vid_index.properties + past_vids = self.vid_index.get_all_properties() # only support Faiss # TODO: We should build vid vector index separately, especially when the vertices may be very large +<<<<<<< HEAD +======= +>>>>>>> 38dce0b (feat(llm): vector db finished) present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py removed_vids = set(past_vids) - set(present_vids) removed_num = self.vid_index.remove(removed_vids) @@ -65,6 +110,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: added_embeddings = asyncio.run(get_embeddings_parallel(self.embedding, vids_to_process)) log.info("Building vector index for %s vertices...", len(added_vids)) self.vid_index.add(added_embeddings, added_vids) +<<<<<<< HEAD self.vid_index.to_index_file(self.index_dir, self.filename_prefix) else: log.debug("No update vertices to build vector index.") @@ -74,4 +120,10 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: "added_vid_vector_num": len(added_vids), } ) +======= + self.vid_index.save_index_by_name(huge_settings.graph_name, "graph_vids") + else: + log.debug("No update vertices to build vector index.") + context.update({"removed_vid_vector_num": removed_num, "added_vid_vector_num": len(added_vids)}) +>>>>>>> 38dce0b (feat(llm): vector db finished) return context diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py index c9aa1876a..ed3f28a2a 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py @@ -16,18 +16,28 @@ # under the License. +<<<<<<< HEAD import asyncio import os from typing import Dict, Any +======= +from typing import Any, Dict +>>>>>>> 38dce0b (feat(llm): vector db finished) <<<<<<< HEAD from hugegraph_llm.config import huge_settings, resource_path, llm_settings from hugegraph_llm.indices.vector_index import VectorIndex ======= from tqdm import tqdm +<<<<<<< HEAD from hugegraph_llm.config import huge_settings, resource_path from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex >>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) +======= + +from hugegraph_llm.config import huge_settings +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +>>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.utils.embedding_utils import ( get_embeddings_parallel, @@ -38,8 +48,9 @@ class BuildVectorIndex: - def __init__(self, embedding: BaseEmbedding): + def __init__(self, embedding: BaseEmbedding, vector_index: type[VectorStoreBase]): self.embedding = embedding +<<<<<<< HEAD <<<<<<< HEAD self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space @@ -53,6 +64,13 @@ def __init__(self, embedding: BaseEmbedding): self.index_dir = str(os.path.join(resource_path, huge_settings.graph_name, "chunks")) self.vector_index = FaissVectorIndex.from_name(self.index_dir) >>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) +======= + self.vector_index = vector_index.from_name( + embedding.get_embedding_dim(), + huge_settings.graph_name, + "chunks", + ) +>>>>>>> 38dce0b (feat(llm): vector db finished) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: if "chunks" not in context: @@ -64,5 +82,9 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: chunks_embedding = asyncio.run(get_embeddings_parallel(self.embedding, chunks)) if len(chunks_embedding) > 0: self.vector_index.add(chunks_embedding, chunks) +<<<<<<< HEAD self.vector_index.to_index_file(self.index_dir, self.filename_prefix) +======= + self.vector_index.save_index_by_name(huge_settings.graph_name, "chunks") +>>>>>>> 38dce0b (feat(llm): vector db finished) return context diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index 89861ead2..1e35841c8 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -18,15 +18,20 @@ import asyncio import os -from typing import Dict, Any, List, Optional +from typing import Any, Dict, List, Optional import pandas as pd +<<<<<<< HEAD <<<<<<< HEAD from hugegraph_llm.config import resource_path, llm_settings, huge_settings from hugegraph_llm.indices.vector_index import VectorIndex, INDEX_FILE_NAME, PROPERTIES_FILE_NAME ======= from hugegraph_llm.config import resource_path, huge_settings, llm_settings +======= +from hugegraph_llm.config import resource_path +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +>>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex >>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) from hugegraph_llm.models.embeddings.base import BaseEmbedding @@ -40,9 +45,12 @@ class GremlinExampleIndexQuery: - def __init__(self, embedding: Optional[BaseEmbedding] = None, num_examples: int = 1): + def __init__( + self, vector_index: type[VectorStoreBase], embedding: Optional[BaseEmbedding] = None, num_examples: int = 1 + ): self.embedding = embedding or Embeddings().get_embedding() self.num_examples = num_examples +<<<<<<< HEAD self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space ) @@ -66,8 +74,15 @@ def _ensure_index_exists(self): os.path.exists(os.path.join(self.index_dir, index_name)) and os.path.exists(os.path.join(self.index_dir, props_name)) ): +======= + if not vector_index.exist("gremlin_examples"): +>>>>>>> 38dce0b (feat(llm): vector db finished) log.warning("No gremlin example index found, will generate one.") + self.vector_index = vector_index.from_name(self.embedding.get_embedding_dim(), "gremlin_examples") + self._build_default_example_index() + else: + self.vector_index = vector_index.from_name(self.embedding.get_embedding_dim(), "gremlin_examples") def _get_match_result(self, context: Dict[str, Any], query: str) -> List[Dict[str, Any]]: if self.num_examples <= 0: @@ -91,7 +106,11 @@ def _build_default_example_index(self): properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict(orient="records") from concurrent.futures import ThreadPoolExecutor +<<<<<<< HEAD # TODO: use asyncio for IO tasks +======= + # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) +>>>>>>> 38dce0b (feat(llm): vector db finished) with ThreadPoolExecutor() as executor: embeddings = list( tqdm( @@ -99,10 +118,15 @@ def _build_default_example_index(self): total=len(properties), ) ) +<<<<<<< HEAD vector_index = FaissVectorIndex(len(embeddings[0])) >>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) vector_index.add(embeddings, properties) vector_index.to_index_file(self.index_dir, self.filename_prefix) +======= + self.vector_index.add(embeddings, properties) + self.vector_index.save_index_by_name("gremlin_examples") +>>>>>>> 38dce0b (feat(llm): vector db finished) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query = context.get("query") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py index 97b26c2c2..929799c67 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py @@ -17,8 +17,9 @@ import os -from typing import Dict, Any, Literal, List, Tuple +from typing import Any, Dict, List, Literal, Tuple +<<<<<<< HEAD <<<<<<< HEAD from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index import VectorIndex @@ -26,10 +27,15 @@ from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex >>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) +======= +from pyhugegraph.client import PyHugeClient + +from hugegraph_llm.config import huge_settings, resource_path +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +>>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.utils.embedding_utils import get_filename_prefix, get_index_folder_name from hugegraph_llm.utils.log import log -from pyhugegraph.client import PyHugeClient class SemanticIdQuery: @@ -38,11 +44,13 @@ class SemanticIdQuery: def __init__( self, embedding: BaseEmbedding, + vector_index: type[VectorStoreBase], by: Literal["query", "keywords"] = "keywords", topk_per_query: int = 10, topk_per_keyword: int = huge_settings.topk_per_keyword, vector_dis_threshold: float = huge_settings.vector_dis_threshold, ): +<<<<<<< HEAD <<<<<<< HEAD self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space @@ -62,6 +70,12 @@ def __init__( ) self.vector_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) >>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) +======= + self.index_dir = str(os.path.join(resource_path, huge_settings.graph_name, "graph_vids")) + self.vector_index = vector_index.from_name( + embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids" + ) +>>>>>>> 38dce0b (feat(llm): vector db finished) self.embedding = embedding self.by = by self.topk_per_query = topk_per_query diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py index c34007d1e..ab5a93912 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py @@ -16,9 +16,9 @@ # under the License. -import os -from typing import Dict, Any +from typing import Any, Dict +<<<<<<< HEAD <<<<<<< HEAD from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index import VectorIndex @@ -26,15 +26,20 @@ from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex >>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) +======= +from hugegraph_llm.config import huge_settings +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +>>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.utils.embedding_utils import get_filename_prefix, get_index_folder_name from hugegraph_llm.utils.log import log class VectorIndexQuery: - def __init__(self, embedding: BaseEmbedding, topk: int = 3): + def __init__(self, vector_index: type[VectorStoreBase], embedding: BaseEmbedding, topk: int = 3): self.embedding = embedding self.topk = topk +<<<<<<< HEAD <<<<<<< HEAD self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space @@ -44,6 +49,9 @@ def __init__(self, embedding: BaseEmbedding, topk: int = 3): llm_settings.embedding_type, getattr(embedding, "model_name", None) ) self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) +======= + self.vector_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "chunks") +>>>>>>> 38dce0b (feat(llm): vector db finished) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query = context.get("query") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py index 7cea5f75a..0443b15c8 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py @@ -16,8 +16,11 @@ # under the License. -from typing import Dict, Any, Optional, Literal, Union, List +from typing import Any, Dict, List, Literal, Optional, Union +from pyhugegraph.client import PyHugeClient + +from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.operators.common_op.check_schema import CheckSchema @@ -31,9 +34,13 @@ from hugegraph_llm.operators.llm_op.disambiguate_data import DisambiguateData from hugegraph_llm.operators.llm_op.info_extract import InfoExtract from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract +<<<<<<< HEAD from hugegraph_llm.operators.llm_op.schema_build import SchemaBuilder from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm from pyhugegraph.client import PyHugeClient +======= +from hugegraph_llm.utils.decorators import log_operator_time, log_time, record_rpm +>>>>>>> 38dce0b (feat(llm): vector db finished) class KgBuilder: @@ -98,14 +105,14 @@ def commit_to_hugegraph(self): self.operators.append(Commit2Graph()) return self - def build_vertex_id_semantic_index(self): + def build_vertex_id_semantic_index(self, vector_index: type[VectorStoreBase]): assert self.embedding - self.operators.append(BuildSemanticIndex(self.embedding)) + self.operators.append(BuildSemanticIndex(self.embedding, vector_index)) return self - def build_vector_index(self): + def build_vector_index(self, vector_index: type[VectorStoreBase]): assert self.embedding - self.operators.append(BuildVectorIndex(self.embedding)) + self.operators.append(BuildVectorIndex(self.embedding, vector_index)) return self def print_result(self): diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index 650834300..edea66f26 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -18,7 +18,7 @@ import asyncio import json import re -from typing import Optional, List, Dict, Any, Union +from typing import Any, Dict, List, Optional, Union from hugegraph_llm.config import prompt from hugegraph_llm.models.llms.base import BaseLLM @@ -29,7 +29,7 @@ class GremlinGenerateSynthesize: def __init__( self, - llm: BaseLLM = None, + llm: BaseLLM | None = None, schema: Optional[Union[dict, str]] = None, vertices: Optional[List[str]] = None, gremlin_prompt: Optional[str] = None, @@ -53,10 +53,14 @@ def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional return None example_strings = [] for example in examples: +<<<<<<< HEAD example_strings.append( f"- query: {example['query']}\n" f"- gremlin:\n```gremlin\n{example['gremlin']}\n```" ) +======= + example_strings.append(f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```") +>>>>>>> 38dce0b (feat(llm): vector db finished) return "\n\n".join(example_strings) def _format_vertices(self, vertices: Optional[List[str]]) -> Optional[str]: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 9645d893e..70c94b7bf 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -19,11 +19,12 @@ import json import os import traceback -from typing import Dict, Any, Union, Optional +from typing import Any, Dict, Optional, Union import gradio as gr from hugegraph_llm.flows.scheduler import SchedulerSingleton +<<<<<<< HEAD from .embedding_utils import get_filename_prefix, get_index_folder_name from .hugegraph_utils import get_hg_client, clean_hg_data from .log import log @@ -33,16 +34,23 @@ from ..indices.vector_index.faiss_vector_store import FaissVectorIndex ======= from ..config import resource_path, huge_settings +======= +from ..config import huge_settings, index_settings, resource_path +>>>>>>> 38dce0b (feat(llm): vector db finished) from ..indices.vector_index.faiss_vector_store import FaissVectorIndex >>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) from ..models.embeddings.init_embedding import Embeddings from ..models.llms.init_llm import LLMs from ..operators.kg_construction_task import KgBuilder +from .hugegraph_utils import clean_hg_data, get_hg_client +from .log import log +from .vector_index_utils import get_vector_index_class, read_documents def get_graph_index_info(): builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) graph_summary_info = builder.fetch_graph_data().run() +<<<<<<< HEAD folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None) @@ -50,10 +58,17 @@ def get_graph_index_info(): vector_index = FaissVectorIndex.from_index_file( str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix, record_miss=False ) +======= + vector_index = get_vector_index_class(index_settings.now_vector_index) + vector_index_entity = vector_index.from_name( + Embeddings().get_embedding().get_embedding_dim(), huge_settings.graph_name, "chunks" + ) + vector_index_info = vector_index_entity.get_vector_index_info() +>>>>>>> 38dce0b (feat(llm): vector db finished) graph_summary_info["vid_index"] = { - "embed_dim": vector_index.index.d, - "num_vectors": vector_index.index.ntotal, - "num_vids": len(vector_index.properties), + "embed_dim": vector_index_info['embed_dim'], + "num_vectors": vector_index_info['vector_info']['chunk_vector_num'], + "num_vids": vector_index_info['vector_info']['graph_properties_vector_num'], } return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) @@ -144,7 +159,14 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: def update_vid_embedding(): +<<<<<<< HEAD scheduler = SchedulerSingleton.get_instance() +======= + vector_index = get_vector_index_class(index_settings.now_vector_index) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder.fetch_graph_data().build_vertex_id_semantic_index(vector_index) + log.debug("Operators: %s", builder.operators) +>>>>>>> 38dce0b (feat(llm): vector db finished) try: return scheduler.schedule_flow("update_vid_embeddings") except Exception as e: # pylint: disable=broad-exception-caught diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index f7610ce9a..7b889d970 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -16,11 +16,12 @@ # under the License. import json -import os +from typing import Type import docx import gradio as gr +<<<<<<< HEAD from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex from hugegraph_llm.models.embeddings.init_embedding import Embeddings, model_map @@ -29,6 +30,14 @@ get_filename_prefix, get_index_folder_name, ) +======= +from hugegraph_llm.config import huge_settings, index_settings +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex +from hugegraph_llm.indices.vector_index.milvus_vector_store import MilvusVectorIndex +from hugegraph_llm.indices.vector_index.qdrant_vector_store import QdrantVectorIndex +from hugegraph_llm.models.embeddings.init_embedding import Embeddings +>>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.models.llms.init_llm import LLMs from hugegraph_llm.operators.kg_construction_task import KgBuilder from hugegraph_llm.utils.hugegraph_utils import get_hg_client @@ -63,6 +72,7 @@ def read_documents(input_file, input_text): # pylint: disable=C0301 def get_vector_index_info(): +<<<<<<< HEAD folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) @@ -74,33 +84,55 @@ def get_vector_index_info(): ) graph_vid_vector_index = FaissVectorIndex.from_index_file( str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix +======= + vector_index = get_vector_index_class(index_settings.now_vector_index) + vector_index_entity = vector_index.from_name( + Embeddings().get_embedding().get_embedding_dim(), huge_settings.graph_name, "chunks" +>>>>>>> 38dce0b (feat(llm): vector db finished) ) + return json.dumps( - { - "embed_dim": chunk_vector_index.index.d, - "vector_info": { - "chunk_vector_num": chunk_vector_index.index.ntotal, - "graph_vid_vector_num": graph_vid_vector_index.index.ntotal, - "graph_properties_vector_num": len(chunk_vector_index.properties), - }, - }, + vector_index_entity.get_vector_index_info(), ensure_ascii=False, indent=2, ) def clean_vector_index(): +<<<<<<< HEAD folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) ) FaissVectorIndex.clean(str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix) +======= + vector_index = get_vector_index_class(index_settings.now_vector_index) + vector_index.clean(huge_settings.graph_name, "chunks") +>>>>>>> 38dce0b (feat(llm): vector db finished) gr.Info("Clean vector index successfully!") def build_vector_index(input_file, input_text): + vector_index = get_vector_index_class(index_settings.now_vector_index) if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) +<<<<<<< HEAD scheduler = SchedulerSingleton.get_instance() return scheduler.schedule_flow("build_vector_index", texts) +======= + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + context = builder.chunk_split(texts, "paragraph", "zh").build_vector_index(vector_index).run() + return json.dumps(context, ensure_ascii=False, indent=2) + + +def get_vector_index_class(vector_index_str: str) -> Type[VectorStoreBase]: + mapping = { + "Faiss": FaissVectorIndex, + "Milvus": MilvusVectorIndex, + "Qdrant": QdrantVectorIndex, + } + ret = mapping.get(vector_index_str) + assert ret + return ret # type: ignore +>>>>>>> 38dce0b (feat(llm): vector db finished) diff --git a/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py index 27c8f2995..128d43b0b 100644 --- a/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py @@ -63,7 +63,7 @@ def test_save_and_load(self): index = MilvusVectorIndex.from_name(test_name) index.add(data_embedding, data) - index.to_index_file(test_name) + index.save_index_by_name(test_name) loaded_index = MilvusVectorIndex.from_name(test_name) diff --git a/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py index 0de26528d..57005fe2b 100644 --- a/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py @@ -64,7 +64,7 @@ def test_save_and_load(self): index = QdrantVectorIndex.from_name(self.name) index.add(data_embedding, data) - index.to_index_file(self.name) + index.save_index_by_name(self.name) loaded_index = QdrantVectorIndex.from_name(self.name) From 9f7d64f9e3cf2db1825ce52849307d05d766bb55 Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Thu, 22 May 2025 09:48:46 +0800 Subject: [PATCH 12/71] feat(llm): updata llm --- hugegraph-llm/src/hugegraph_llm/config/__init__.py | 4 ++-- hugegraph-llm/src/hugegraph_llm/config/generate.py | 3 ++- hugegraph-llm/src/hugegraph_llm/config/llm_config.py | 6 +++--- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/config/__init__.py b/hugegraph-llm/src/hugegraph_llm/config/__init__.py index 7ff7f2938..f7f9cf290 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/__init__.py +++ b/hugegraph-llm/src/hugegraph_llm/config/__init__.py @@ -22,10 +22,10 @@ from hugegraph_llm.config.index_config import IndexConfig -from .prompt_config import PromptConfig -from .hugegraph_config import HugeGraphConfig from .admin_config import AdminConfig +from .hugegraph_config import HugeGraphConfig from .llm_config import LLMConfig +from .prompt_config import PromptConfig llm_settings = LLMConfig() prompt = PromptConfig(llm_settings) diff --git a/hugegraph-llm/src/hugegraph_llm/config/generate.py b/hugegraph-llm/src/hugegraph_llm/config/generate.py index 4b40e899f..9574b7b06 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/generate.py +++ b/hugegraph-llm/src/hugegraph_llm/config/generate.py @@ -18,7 +18,7 @@ import argparse -from hugegraph_llm.config import huge_settings, admin_settings, llm_settings, PromptConfig +from hugegraph_llm.config import PromptConfig, admin_settings, huge_settings, index_settings, llm_settings if __name__ == "__main__": parser = argparse.ArgumentParser(description="Generate hugegraph-llm config file") @@ -30,4 +30,5 @@ huge_settings.generate_env() admin_settings.generate_env() llm_settings.generate_env() + index_settings.generate_env() PromptConfig(llm_settings).generate_yaml_file() diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index e3e75a66c..493eca287 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -62,9 +62,9 @@ class LLMConfig(BaseConfig): ollama_text2gql_host: str = "127.0.0.1" ollama_text2gql_port: int = 11434 ollama_text2gql_language_model: str | None = None - ollama_embedding_host: str = os.getenv("OLLAMA_EMBEDDING_HOST", "127.0.0.1") - ollama_embedding_port: int = int(os.getenv("OLLAMA_EMBEDDING_PORT", 11434)) - ollama_embedding_model: str = os.getenv("OLLAMA_EMBEDDING_MODEL", 'quentinz/bge-large-zh-v1.5') + ollama_embedding_host: str = "127.0.0.1" + ollama_embedding_port: int = int(11434) + ollama_embedding_model: str = 'quentinz/bge-large-zh-v1.5' ollama_embedding_model_dim: Optional[int] = ( int(os.getenv("OLLAMA_EMBEDDING_MODEL_DIM")) if os.getenv("OLLAMA_EMBEDDING_MODEL_DIM") else None # type:ignore ) From f470605303e67e7a74c9f0401e2c3e1972fb88c6 Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Thu, 22 May 2025 09:56:14 +0800 Subject: [PATCH 13/71] feat(llm): nexpected-keyword-arg,unused-import --- .../demo/rag_demo/text2gremlin_block.py | 4 ++ .../src/hugegraph_llm/indices/graph_index.py | 1 + .../hugegraph_llm/operators/graph_rag_task.py | 3 +- .../index_op/gremlin_example_index_query.py | 58 +------------------ .../hugegraph_llm/utils/vector_index_utils.py | 29 +--------- 5 files changed, 8 insertions(+), 87 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 11103d7c4..7ebed651a 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -388,9 +388,13 @@ def graph_rag_recall( rag = RAGPipeline() rag.extract_keywords().keywords_to_vid( <<<<<<< HEAD +<<<<<<< HEAD ======= vector_index=index_settings.now_vector_index, >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + vector_index_str=index_settings.now_vector_index, +>>>>>>> dd3b085 (feat(llm): nexpected-keyword-arg,unused-import) vector_dis_threshold=vector_dis_threshold, topk_per_keyword=topk_per_keyword, ) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py b/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py index 694ca014d..31b3d21f3 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py @@ -17,6 +17,7 @@ from typing import Optional + from pyhugegraph.client import PyHugeClient from hugegraph_llm.config import huge_settings diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index 80da2b57e..b37707ca3 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -18,8 +18,7 @@ from typing import Any, Dict, List, Literal, Optional -from hugegraph_llm.config import huge_settings, index_settings, prompt -from hugegraph_llm.indices.vector_index.base import VectorStoreBase +from hugegraph_llm.config import huge_settings, prompt from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.models.llms.base import BaseLLM diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index 1e35841c8..28baca370 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -22,18 +22,8 @@ import pandas as pd -<<<<<<< HEAD -<<<<<<< HEAD -from hugegraph_llm.config import resource_path, llm_settings, huge_settings -from hugegraph_llm.indices.vector_index import VectorIndex, INDEX_FILE_NAME, PROPERTIES_FILE_NAME -======= -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -======= -from hugegraph_llm.config import resource_path +from hugegraph_llm.config import resource_path, huge_settings from hugegraph_llm.indices.vector_index.base import VectorStoreBase ->>>>>>> 38dce0b (feat(llm): vector db finished) -from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.utils.embedding_utils import ( @@ -50,33 +40,7 @@ def __init__( ): self.embedding = embedding or Embeddings().get_embedding() self.num_examples = num_examples -<<<<<<< HEAD - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "gremlin_examples")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(self.embedding, "model_name", None) - ) - self._ensure_index_exists() - self.vector_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) - - def _ensure_index_exists(self): - index_name = ( - f"{self.filename_prefix}_{INDEX_FILE_NAME}" if self.filename_prefix else INDEX_FILE_NAME - ) - props_name = ( - f"{self.filename_prefix}_{PROPERTIES_FILE_NAME}" - if self.filename_prefix - else PROPERTIES_FILE_NAME - ) - if not ( - os.path.exists(os.path.join(self.index_dir, index_name)) - and os.path.exists(os.path.join(self.index_dir, props_name)) - ): -======= if not vector_index.exist("gremlin_examples"): ->>>>>>> 38dce0b (feat(llm): vector db finished) log.warning("No gremlin example index found, will generate one.") self.vector_index = vector_index.from_name(self.embedding.get_embedding_dim(), "gremlin_examples") @@ -94,23 +58,10 @@ def _get_match_result(self, context: Dict[str, Any], query: str) -> List[Dict[st return self.vector_index.search(query_embedding, self.num_examples, dis_threshold=1.8) def _build_default_example_index(self): -<<<<<<< HEAD - properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict( - orient="records" - ) - # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) - queries = [row["query"] for row in properties] - embeddings = asyncio.run(get_embeddings_parallel(self.embedding, queries)) - vector_index = FaissVectorIndex(len(embeddings[0])) -======= properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict(orient="records") from concurrent.futures import ThreadPoolExecutor -<<<<<<< HEAD - # TODO: use asyncio for IO tasks -======= # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) ->>>>>>> 38dce0b (feat(llm): vector db finished) with ThreadPoolExecutor() as executor: embeddings = list( tqdm( @@ -118,15 +69,8 @@ def _build_default_example_index(self): total=len(properties), ) ) -<<<<<<< HEAD - vector_index = FaissVectorIndex(len(embeddings[0])) ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) - vector_index.add(embeddings, properties) - vector_index.to_index_file(self.index_dir, self.filename_prefix) -======= self.vector_index.add(embeddings, properties) self.vector_index.save_index_by_name("gremlin_examples") ->>>>>>> 38dce0b (feat(llm): vector db finished) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query = context.get("query") diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 7b889d970..d8c87fdd6 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -72,43 +72,21 @@ def read_documents(input_file, input_text): # pylint: disable=C0301 def get_vector_index_info(): -<<<<<<< HEAD - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) - ) - chunk_vector_index = FaissVectorIndex.from_index_file( - str(os.path.join(resource_path, folder_name, "chunks")), - filename_prefix, - record_miss=False, - ) - graph_vid_vector_index = FaissVectorIndex.from_index_file( - str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix -======= vector_index = get_vector_index_class(index_settings.now_vector_index) vector_index_entity = vector_index.from_name( Embeddings().get_embedding().get_embedding_dim(), huge_settings.graph_name, "chunks" ->>>>>>> 38dce0b (feat(llm): vector db finished) ) return json.dumps( - vector_index_entity.get_vector_index_info(), + {**vector_index_entity.get_vector_index_info(), 'now_vector_index': index_settings.now_vector_index}, ensure_ascii=False, indent=2, ) def clean_vector_index(): -<<<<<<< HEAD - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) - ) - FaissVectorIndex.clean(str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix) -======= vector_index = get_vector_index_class(index_settings.now_vector_index) vector_index.clean(huge_settings.graph_name, "chunks") ->>>>>>> 38dce0b (feat(llm): vector db finished) gr.Info("Clean vector index successfully!") @@ -117,10 +95,6 @@ def build_vector_index(input_file, input_text): if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) -<<<<<<< HEAD - scheduler = SchedulerSingleton.get_instance() - return scheduler.schedule_flow("build_vector_index", texts) -======= builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) context = builder.chunk_split(texts, "paragraph", "zh").build_vector_index(vector_index).run() return json.dumps(context, ensure_ascii=False, indent=2) @@ -135,4 +109,3 @@ def get_vector_index_class(vector_index_str: str) -> Type[VectorStoreBase]: ret = mapping.get(vector_index_str) assert ret return ret # type: ignore ->>>>>>> 38dce0b (feat(llm): vector db finished) From fce80a9fe58add13e8a1ae9b1fb90559d1e4ff00 Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Thu, 22 May 2025 10:02:54 +0800 Subject: [PATCH 14/71] feat(llm): fit unitest --- .../src/tests/indices/test_milvus_vector_index.py | 8 ++++---- .../src/tests/indices/test_qdrant_vector_index.py | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py index 128d43b0b..b1ac0f209 100644 --- a/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py @@ -39,7 +39,7 @@ def test_vector_index(self): ] data_embedding = [embedder.get_text_embedding(d) for d in data] - index = MilvusVectorIndex.from_name(test_name) + index = MilvusVectorIndex.from_name(1024, test_name) index.add(data_embedding, data) query = "腾讯的合伙人有哪些？" @@ -60,12 +60,12 @@ def test_save_and_load(self): ] data_embedding = [embedder.get_text_embedding(d) for d in data] - index = MilvusVectorIndex.from_name(test_name) + index = MilvusVectorIndex.from_name(1024, test_name) index.add(data_embedding, data) index.save_index_by_name(test_name) - loaded_index = MilvusVectorIndex.from_name(test_name) + loaded_index = MilvusVectorIndex.from_name(1024, test_name) query = "腾讯的合伙人有哪些？" query_vector = embedder.get_text_embedding(query) @@ -83,7 +83,7 @@ def test_remove_entries(self): ] data_embedding = [embedder.get_text_embedding(d) for d in data] - index = MilvusVectorIndex.from_name(test_name) + index = MilvusVectorIndex.from_name(1024, test_name) index.add(data_embedding, data) query = "合伙人" diff --git a/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py index 57005fe2b..1e0768051 100644 --- a/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py @@ -40,7 +40,7 @@ def test_vector_index(self): ] data_embedding = [embedder.get_text_embedding(d) for d in data] - index = QdrantVectorIndex.from_name(self.name) + index = QdrantVectorIndex.from_name(1024, self.name) index.add(data_embedding, data) query = "腾讯的合伙人有哪些？" @@ -61,12 +61,12 @@ def test_save_and_load(self): ] data_embedding = [embedder.get_text_embedding(d) for d in data] - index = QdrantVectorIndex.from_name(self.name) + index = QdrantVectorIndex.from_name(1024, self.name) index.add(data_embedding, data) index.save_index_by_name(self.name) - loaded_index = QdrantVectorIndex.from_name(self.name) + loaded_index = QdrantVectorIndex.from_name(1024, self.name) query = "腾讯的合伙人有哪些？" query_vector = embedder.get_text_embedding(query) @@ -85,7 +85,7 @@ def test_remove_entries(self): ] data_embedding = [embedder.get_text_embedding(d) for d in data] - index = QdrantVectorIndex.from_name(self.name) + index = QdrantVectorIndex.from_name(1024, self.name) index.add(data_embedding, data) query = "合伙人" From f3a8a260182f287d34cadaea1aca0262928f2aa1 Mon Sep 17 00:00:00 2001 From: mikumifa <1055069518@qq.com> Date: Thu, 22 May 2025 10:11:31 +0800 Subject: [PATCH 15/71] feat(llm): use lambda --- .../demo/rag_demo/configs_block.py | 214 ++++++++++++++---- .../vector_index/faiss_vector_store.py | 90 -------- 2 files changed, 176 insertions(+), 128 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index e2633229e..01fc2c2da 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -303,11 +303,12 @@ def create_configs_block() -> list: with gr.Row(): graph_config_input = [ gr.Textbox( - value=huge_settings.graph_url, + value=lambda: huge_settings.graph_url, label="url", info="IP:PORT (e.g. 127.0.0.1:8080) or full URL (e.g. http://127.0.0.1:8080)", ), gr.Textbox( +<<<<<<< HEAD <<<<<<< HEAD value=huge_settings.graph_name, label="graph", @@ -325,14 +326,26 @@ def create_configs_block() -> list: info="Password for graph server auth", ======= value=huge_settings.graph_name, label="graph", info="The graph name of HugeGraph-Server instance" +======= + value=lambda: huge_settings.graph_name, + label="graph", + info="The graph name of HugeGraph-Server instance", +>>>>>>> f42fa9b (feat(llm): use lambda) ), - gr.Textbox(value=huge_settings.graph_user, label="user", info="Username for graph server auth"), + gr.Textbox(value=lambda: huge_settings.graph_user, label="user", info="Username for graph server auth"), gr.Textbox( +<<<<<<< HEAD value=huge_settings.graph_pwd, label="pwd", type="password", info="Password for graph server auth" >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + value=lambda: huge_settings.graph_pwd, + label="pwd", + type="password", + info="Password for graph server auth", +>>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=huge_settings.graph_space, + value=lambda: huge_settings.graph_space, label="graphspace (Optional)", info="Namespace for multi-tenant scenarios (leave empty if not using graphspaces)", ), @@ -368,6 +381,7 @@ def chat_llm_settings(llm_type): if llm_type == "openai": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "openai_chat_api_key"), label="api_key", @@ -402,33 +416,51 @@ def chat_llm_settings(llm_type): ), ======= value=getattr(llm_settings, "openai_chat_api_key"), label="api_key", type="password" +======= + value=lambda: getattr(llm_settings, "openai_chat_api_key"), label="api_key", type="password" +>>>>>>> f42fa9b (feat(llm): use lambda) ), - gr.Textbox(value=getattr(llm_settings, "openai_chat_api_base"), label="api_base"), - gr.Textbox(value=getattr(llm_settings, "openai_chat_language_model"), label="model_name"), - gr.Textbox(value=getattr(llm_settings, "openai_chat_tokens"), label="max_token"), + gr.Textbox(value=lambda: getattr(llm_settings, "openai_chat_api_base"), label="api_base"), + gr.Textbox( + value=lambda: getattr(llm_settings, "openai_chat_language_model"), label="model_name" + ), + gr.Textbox(value=lambda: getattr(llm_settings, "openai_chat_tokens"), label="max_token"), ] elif llm_type == "ollama/local": llm_config_input = [ - gr.Textbox(value=getattr(llm_settings, "ollama_chat_host"), label="host"), - gr.Textbox(value=str(getattr(llm_settings, "ollama_chat_port")), label="port"), - gr.Textbox(value=getattr(llm_settings, "ollama_chat_language_model"), label="model_name"), + gr.Textbox(value=lambda: getattr(llm_settings, "ollama_chat_host"), label="host"), + gr.Textbox(value=lambda: str(getattr(llm_settings, "ollama_chat_port")), label="port"), + gr.Textbox( + value=lambda: getattr(llm_settings, "ollama_chat_language_model"), label="model_name" + ), gr.Textbox(value="", visible=False), ] elif llm_type == "qianfan_wenxin": llm_config_input = [ gr.Textbox( - value=getattr(llm_settings, "qianfan_chat_api_key"), label="api_key", type="password" + value=lambda: getattr(llm_settings, "qianfan_chat_api_key"), + label="api_key", + type="password", ), gr.Textbox( - value=getattr(llm_settings, "qianfan_chat_secret_key"), label="secret_key", type="password" + value=lambda: getattr(llm_settings, "qianfan_chat_secret_key"), + label="secret_key", + type="password", ), + gr.Textbox( + value=lambda: getattr(llm_settings, "qianfan_chat_language_model"), label="model_name" + ), +<<<<<<< HEAD gr.Textbox(value=getattr(llm_settings, "qianfan_chat_language_model"), label="model_name"), >>>>>>> 38dce0b (feat(llm): vector db finished) +======= +>>>>>>> f42fa9b (feat(llm): use lambda) gr.Textbox(value="", visible=False), ] elif llm_type == "litellm": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "litellm_chat_api_key"), label="api_key", @@ -436,17 +468,23 @@ def chat_llm_settings(llm_type): ======= value=getattr(llm_settings, "litellm_chat_api_key"), label="api_key", type="password" >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + value=lambda: getattr(llm_settings, "litellm_chat_api_key"), + label="api_key", + type="password", +>>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=getattr(llm_settings, "litellm_chat_api_base"), + value=lambda: getattr(llm_settings, "litellm_chat_api_base"), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr(llm_settings, "litellm_chat_language_model"), + value=lambda: getattr(llm_settings, "litellm_chat_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), +<<<<<<< HEAD <<<<<<< HEAD gr.Textbox( value=getattr(llm_settings, "litellm_chat_tokens"), @@ -455,6 +493,9 @@ def chat_llm_settings(llm_type): ======= gr.Textbox(value=getattr(llm_settings, "litellm_chat_tokens"), label="max_token"), >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + gr.Textbox(value=lambda: getattr(llm_settings, "litellm_chat_tokens"), label="max_token"), +>>>>>>> f42fa9b (feat(llm): use lambda) ] else: llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] @@ -498,6 +539,7 @@ def extract_llm_settings(llm_type): if llm_type == "openai": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "openai_extract_api_key"), label="api_key", @@ -532,35 +574,53 @@ def extract_llm_settings(llm_type): ), ======= value=getattr(llm_settings, "openai_extract_api_key"), label="api_key", type="password" +======= + value=lambda: getattr(llm_settings, "openai_extract_api_key"), + label="api_key", + type="password", +>>>>>>> f42fa9b (feat(llm): use lambda) ), - gr.Textbox(value=getattr(llm_settings, "openai_extract_api_base"), label="api_base"), - gr.Textbox(value=getattr(llm_settings, "openai_extract_language_model"), label="model_name"), - gr.Textbox(value=getattr(llm_settings, "openai_extract_tokens"), label="max_token"), + gr.Textbox(value=lambda: getattr(llm_settings, "openai_extract_api_base"), label="api_base"), + gr.Textbox( + value=lambda: getattr(llm_settings, "openai_extract_language_model"), label="model_name" + ), + gr.Textbox(value=lambda: getattr(llm_settings, "openai_extract_tokens"), label="max_token"), ] elif llm_type == "ollama/local": llm_config_input = [ - gr.Textbox(value=getattr(llm_settings, "ollama_extract_host"), label="host"), - gr.Textbox(value=str(getattr(llm_settings, "ollama_extract_port")), label="port"), - gr.Textbox(value=getattr(llm_settings, "ollama_extract_language_model"), label="model_name"), + gr.Textbox(value=lambda: getattr(llm_settings, "ollama_extract_host"), label="host"), + gr.Textbox(value=lambda: str(getattr(llm_settings, "ollama_extract_port")), label="port"), + gr.Textbox( + value=lambda: getattr(llm_settings, "ollama_extract_language_model"), label="model_name" + ), gr.Textbox(value="", visible=False), ] elif llm_type == "qianfan_wenxin": llm_config_input = [ gr.Textbox( - value=getattr(llm_settings, "qianfan_extract_api_key"), label="api_key", type="password" + value=lambda: getattr(llm_settings, "qianfan_extract_api_key"), + label="api_key", + type="password", ), gr.Textbox( - value=getattr(llm_settings, "qianfan_extract_secret_key"), + value=lambda: getattr(llm_settings, "qianfan_extract_secret_key"), label="secret_key", type="password", ), +<<<<<<< HEAD gr.Textbox(value=getattr(llm_settings, "qianfan_extract_language_model"), label="model_name"), >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + gr.Textbox( + value=lambda: getattr(llm_settings, "qianfan_extract_language_model"), label="model_name" + ), +>>>>>>> f42fa9b (feat(llm): use lambda) gr.Textbox(value="", visible=False), ] elif llm_type == "litellm": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "litellm_extract_api_key"), label="api_key", @@ -568,17 +628,23 @@ def extract_llm_settings(llm_type): ======= value=getattr(llm_settings, "litellm_extract_api_key"), label="api_key", type="password" >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + value=lambda: getattr(llm_settings, "litellm_extract_api_key"), + label="api_key", + type="password", +>>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=getattr(llm_settings, "litellm_extract_api_base"), + value=lambda: getattr(llm_settings, "litellm_extract_api_base"), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr(llm_settings, "litellm_extract_language_model"), + value=lambda: getattr(llm_settings, "litellm_extract_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), +<<<<<<< HEAD <<<<<<< HEAD gr.Textbox( value=getattr(llm_settings, "litellm_extract_tokens"), @@ -587,6 +653,9 @@ def extract_llm_settings(llm_type): ======= gr.Textbox(value=getattr(llm_settings, "litellm_extract_tokens"), label="max_token"), >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + gr.Textbox(value=lambda: getattr(llm_settings, "litellm_extract_tokens"), label="max_token"), +>>>>>>> f42fa9b (feat(llm): use lambda) ] else: llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] @@ -613,6 +682,7 @@ def text2gql_llm_settings(llm_type): if llm_type == "openai": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "openai_text2gql_api_key"), label="api_key", @@ -647,35 +717,53 @@ def text2gql_llm_settings(llm_type): ), ======= value=getattr(llm_settings, "openai_text2gql_api_key"), label="api_key", type="password" +======= + value=lambda: getattr(llm_settings, "openai_text2gql_api_key"), + label="api_key", + type="password", +>>>>>>> f42fa9b (feat(llm): use lambda) + ), + gr.Textbox(value=lambda: getattr(llm_settings, "openai_text2gql_api_base"), label="api_base"), + gr.Textbox( + value=lambda: getattr(llm_settings, "openai_text2gql_language_model"), label="model_name" ), - gr.Textbox(value=getattr(llm_settings, "openai_text2gql_api_base"), label="api_base"), - gr.Textbox(value=getattr(llm_settings, "openai_text2gql_language_model"), label="model_name"), - gr.Textbox(value=getattr(llm_settings, "openai_text2gql_tokens"), label="max_token"), + gr.Textbox(value=lambda: getattr(llm_settings, "openai_text2gql_tokens"), label="max_token"), ] elif llm_type == "ollama/local": llm_config_input = [ - gr.Textbox(value=getattr(llm_settings, "ollama_text2gql_host"), label="host"), - gr.Textbox(value=str(getattr(llm_settings, "ollama_text2gql_port")), label="port"), - gr.Textbox(value=getattr(llm_settings, "ollama_text2gql_language_model"), label="model_name"), + gr.Textbox(value=lambda: getattr(llm_settings, "ollama_text2gql_host"), label="host"), + gr.Textbox(value=lambda: str(getattr(llm_settings, "ollama_text2gql_port")), label="port"), + gr.Textbox( + value=lambda: getattr(llm_settings, "ollama_text2gql_language_model"), label="model_name" + ), gr.Textbox(value="", visible=False), ] elif llm_type == "qianfan_wenxin": llm_config_input = [ gr.Textbox( - value=getattr(llm_settings, "qianfan_text2gql_api_key"), label="api_key", type="password" + value=lambda: getattr(llm_settings, "qianfan_text2gql_api_key"), + label="api_key", + type="password", ), gr.Textbox( - value=getattr(llm_settings, "qianfan_text2gql_secret_key"), + value=lambda: getattr(llm_settings, "qianfan_text2gql_secret_key"), label="secret_key", type="password", ), +<<<<<<< HEAD gr.Textbox(value=getattr(llm_settings, "qianfan_text2gql_language_model"), label="model_name"), >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + gr.Textbox( + value=lambda: getattr(llm_settings, "qianfan_text2gql_language_model"), label="model_name" + ), +>>>>>>> f42fa9b (feat(llm): use lambda) gr.Textbox(value="", visible=False), ] elif llm_type == "litellm": llm_config_input = [ gr.Textbox( +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "litellm_text2gql_api_key"), label="api_key", @@ -683,17 +771,23 @@ def text2gql_llm_settings(llm_type): ======= value=getattr(llm_settings, "litellm_text2gql_api_key"), label="api_key", type="password" >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + value=lambda: getattr(llm_settings, "litellm_text2gql_api_key"), + label="api_key", + type="password", +>>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=getattr(llm_settings, "litellm_text2gql_api_base"), + value=lambda: getattr(llm_settings, "litellm_text2gql_api_base"), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr(llm_settings, "litellm_text2gql_language_model"), + value=lambda: getattr(llm_settings, "litellm_text2gql_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), +<<<<<<< HEAD <<<<<<< HEAD gr.Textbox( value=getattr(llm_settings, "litellm_text2gql_tokens"), @@ -702,6 +796,9 @@ def text2gql_llm_settings(llm_type): ======= gr.Textbox(value=getattr(llm_settings, "litellm_text2gql_tokens"), label="max_token"), >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + gr.Textbox(value=lambda: getattr(llm_settings, "litellm_text2gql_tokens"), label="max_token"), +>>>>>>> f42fa9b (feat(llm): use lambda) ] else: llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] @@ -725,6 +822,7 @@ def embedding_settings(embedding_type): if embedding_type == "openai": with gr.Row(): embedding_config_input = [ +<<<<<<< HEAD <<<<<<< HEAD gr.Textbox( value=llm_settings.openai_embedding_api_key, @@ -745,10 +843,19 @@ def embedding_settings(embedding_type): gr.Textbox(value=llm_settings.openai_embedding_model, label="model_name"), gr.Textbox(value=str(llm_settings.openai_embedding_model_dim), label="model_dim"), >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + gr.Textbox( + value=lambda: llm_settings.openai_embedding_api_key, label="api_key", type="password" + ), + gr.Textbox(value=lambda: llm_settings.openai_embedding_api_base, label="api_base"), + gr.Textbox(value=lambda: llm_settings.openai_embedding_model, label="model_name"), + gr.Textbox(value=lambda: str(llm_settings.openai_embedding_model_dim), label="model_dim"), +>>>>>>> f42fa9b (feat(llm): use lambda) ] elif embedding_type == "ollama/local": with gr.Row(): embedding_config_input = [ +<<<<<<< HEAD gr.Textbox(value=llm_settings.ollama_embedding_host, label="host"), gr.Textbox(value=str(llm_settings.ollama_embedding_port), label="port"), <<<<<<< HEAD @@ -759,22 +866,36 @@ def embedding_settings(embedding_type): ======= gr.Textbox(value=llm_settings.ollama_embedding_model, label="model_name"), gr.Textbox(value=str(llm_settings.ollama_embedding_model_dim), label="model_dim"), +======= + gr.Textbox(value=lambda: llm_settings.ollama_embedding_host, label="host"), + gr.Textbox(value=lambda: str(llm_settings.ollama_embedding_port), label="port"), + gr.Textbox(value=lambda: llm_settings.ollama_embedding_model, label="model_name"), + gr.Textbox(value=lambda: str(llm_settings.ollama_embedding_model_dim), label="model_dim"), +>>>>>>> f42fa9b (feat(llm): use lambda) ] elif embedding_type == "qianfan_wenxin": with gr.Row(): embedding_config_input = [ - gr.Textbox(value=llm_settings.qianfan_embedding_api_key, label="api_key", type="password"), gr.Textbox( - value=llm_settings.qianfan_embedding_secret_key, label="secret_key", type="password" + value=lambda: llm_settings.qianfan_embedding_api_key, label="api_key", type="password" ), +<<<<<<< HEAD gr.Textbox(value=llm_settings.qianfan_embedding_model, label="model_name"), gr.Textbox(value=str(llm_settings.qianfan_embedding_model_dim), label="model_dim"), >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + gr.Textbox( + value=lambda: llm_settings.qianfan_embedding_secret_key, label="secret_key", type="password" + ), + gr.Textbox(value=lambda: llm_settings.qianfan_embedding_model, label="model_name"), + gr.Textbox(value=lambda: str(llm_settings.qianfan_embedding_model_dim), label="model_dim"), +>>>>>>> f42fa9b (feat(llm): use lambda) ] elif embedding_type == "litellm": with gr.Row(): embedding_config_input = [ gr.Textbox( +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "litellm_embedding_api_key"), label="api_key", @@ -782,21 +903,28 @@ def embedding_settings(embedding_type): ======= value=getattr(llm_settings, "litellm_embedding_api_key"), label="api_key", type="password" >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + value=lambda: getattr(llm_settings, "litellm_embedding_api_key"), + label="api_key", + type="password", +>>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=getattr(llm_settings, "litellm_embedding_api_base"), + value=lambda: getattr(llm_settings, "litellm_embedding_api_base"), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr(llm_settings, "litellm_embedding_model"), + value=lambda: getattr(llm_settings, "litellm_embedding_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/embedding/supported_embedding", ), <<<<<<< HEAD ======= gr.Textbox( - value=getattr(llm_settings, "litellm_embedding_model_dim"), label="model_dim", type="text" + value=lambda: getattr(llm_settings, "litellm_embedding_model_dim"), + label="model_dim", + type="text", ), >>>>>>> 38dce0b (feat(llm): vector db finished) ] @@ -829,6 +957,7 @@ def reranker_settings(reranker_type): if reranker_type == "cohere": with gr.Row(): reranker_config_input = [ +<<<<<<< HEAD gr.Textbox( value=llm_settings.reranker_api_key, label="api_key", @@ -836,15 +965,24 @@ def reranker_settings(reranker_type): ), gr.Textbox(value=llm_settings.reranker_model, label="model"), gr.Textbox(value=llm_settings.cohere_base_url, label="base_url"), +======= + gr.Textbox(value=lambda: llm_settings.reranker_api_key, label="api_key", type="password"), + gr.Textbox(value=lambda: llm_settings.reranker_model, label="model"), + gr.Textbox(value=lambda: llm_settings.cohere_base_url, label="base_url"), +>>>>>>> f42fa9b (feat(llm): use lambda) ] elif reranker_type == "siliconflow": with gr.Row(): reranker_config_input = [ +<<<<<<< HEAD gr.Textbox( value=llm_settings.reranker_api_key, label="api_key", type="password", ), +======= + gr.Textbox(value=lambda: llm_settings.reranker_api_key, label="api_key", type="password"), +>>>>>>> f42fa9b (feat(llm): use lambda) gr.Textbox( value="BAAI/bge-reranker-v2-m3", label="model", diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py index ef731db39..a8f23155a 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py @@ -36,26 +36,10 @@ def __init__(self, embed_dim: int = 1024): self.index = faiss.IndexFlatL2(embed_dim) self.properties: list[Any] = [] -<<<<<<< HEAD - - - def to_index_file(self, dir_path: str, filename_prefix: str = None): - """Save index to files, supporting model-specific filenames.""" - if not os.path.exists(dir_path): - os.makedirs(dir_path) - - index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = ( - f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME - ) - index_file = os.path.join(dir_path, index_name) - properties_file = os.path.join(dir_path, property_name) -======= def save_index_by_name(self, *name: str): os.makedirs(os.path.join(resource_path, *name), exist_ok=True) index_file = os.path.join(resource_path, *name, INDEX_FILE_NAME) properties_file = os.path.join(resource_path, *name, PROPERTIES_FILE_NAME) ->>>>>>> 38dce0b (feat(llm): vector db finished) faiss.write_index(self.index, index_file) with open(properties_file, "wb") as f: pkl.dump(self.properties, f) @@ -121,79 +105,6 @@ def get_vector_index_info( } @staticmethod -<<<<<<< HEAD - def from_index_file( - dir_path: str, filename_prefix: str | None = None, record_miss: bool = True - ) -> "FaissVectorIndex": - """Load index from files, supporting model-specific filenames. - - If prefixed files are missing, optionally warn and return an empty index. - Also validates vector/property count consistency. - """ - index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = ( - f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME - ) - index_file = os.path.join(dir_path, index_name) - properties_file = os.path.join(dir_path, property_name) - missing = [p for p in [index_file, properties_file] if not os.path.exists(p)] - if missing: - if record_miss: - log.warning( - "Missing vector files: %s. Need create a new one for it.", ", ".join(missing) - ) - return FaissVectorIndex() - - try: - faiss_index = faiss.read_index(index_file) - with open(properties_file, "rb") as f: - properties = pkl.load(f) - except (RuntimeError, pkl.UnpicklingError, OSError) as e: # pragma: no cover - log.error( - "Failed to load index files for model '%s': %s", - filename_prefix or "default", - e, - ) - raise RuntimeError( - f"Could not load index files for model '{filename_prefix or 'default'}'. " - f"Original error ({type(e).__name__}): {e}" - ) from e - - if faiss_index.ntotal != len(properties): - raise RuntimeError( - f"Data inconsistency: index for model '{filename_prefix or 'default'}' has " - f"{faiss_index.ntotal} vectors, but {len(properties)} properties." - ) - - embed_dim = faiss_index.d - vector_index = FaissVectorIndex(embed_dim) - vector_index.index = faiss_index - vector_index.properties = properties - return vector_index - - @staticmethod - def clean(dir_path: str, filename_prefix: str = None): - """Clean index files, supporting model-specific filenames. - - This method deletes the index and properties files associated with a specific model. - If model_name is None, it targets the default files. - """ - index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = ( - f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME - ) - index_file = os.path.join(dir_path, index_name) - properties_file = os.path.join(dir_path, property_name) - - for file in [index_file, properties_file]: - if os.path.exists(file): - try: - os.remove(file) - log.info("Removed index file: %s", file) - except OSError as e: - log.error("Error removing file %s: %s", file, e) - -======= def clean(*name: str): index_file = os.path.join(resource_path, *name, INDEX_FILE_NAME) properties_file = os.path.join(resource_path, *name, PROPERTIES_FILE_NAME) @@ -201,7 +112,6 @@ def clean(*name: str): os.remove(index_file) if os.path.exists(properties_file): os.remove(properties_file) ->>>>>>> 38dce0b (feat(llm): vector db finished) @staticmethod def from_name(embed_dim: int, *name: str) -> "FaissVectorIndex": From 77dd3861e65a7a3b3af0511a8321af21fb470a7b Mon Sep 17 00:00:00 2001 From: lingxiao Date: Tue, 12 Aug 2025 10:38:32 +0800 Subject: [PATCH 16/71] style: format code with black line-length 120 --- .../src/hugegraph_llm/api/admin_api.py | 4 +- .../hugegraph_llm/api/models/rag_requests.py | 8 + .../src/hugegraph_llm/config/index_config.py | 2 +- .../demo/rag_demo/admin_block.py | 4 + .../src/hugegraph_llm/demo/rag_demo/app.py | 8 +- .../demo/rag_demo/configs_block.py | 33 ++-- .../demo/rag_demo/other_block.py | 4 + .../hugegraph_llm/demo/rag_demo/rag_block.py | 4 +- .../demo/rag_demo/text2gremlin_block.py | 24 ++- .../demo/rag_demo/vector_graph_block.py | 172 +++++++++++++++++- .../src/hugegraph_llm/document/chunk_split.py | 10 + .../vector_index/milvus_vector_store.py | 6 +- .../vector_index/qdrant_vector_store.py | 6 +- .../hugegraph_llm/middleware/middleware.py | 7 +- .../models/embeddings/init_embedding.py | 4 +- .../hugegraph_llm/models/embeddings/ollama.py | 2 +- .../src/hugegraph_llm/models/llms/ollama.py | 4 + .../operators/document_op/chunk_split.py | 4 +- .../operators/document_op/word_extract.py | 4 +- .../hugegraph_llm/operators/graph_rag_task.py | 4 + .../operators/llm_op/answer_synthesize.py | 74 ++------ .../operators/llm_op/disambiguate_data.py | 4 + .../operators/llm_op/info_extract.py | 21 +++ .../operators/llm_op/keyword_extract.py | 8 +- .../llm_op/property_graph_extract.py | 4 + .../operators/llm_op/schema_build.py | 4 + .../llm_op/unstructured_data_utils.py | 5 +- .../hugegraph_llm/utils/graph_index_utils.py | 43 ++++- .../hugegraph_llm/utils/hugegraph_utils.py | 8 +- hugegraph-llm/src/hugegraph_llm/utils/log.py | 4 + .../hugegraph_llm/utils/vector_index_utils.py | 2 +- .../tests/models/llms/test_ollama_client.py | 4 + .../src/pyhugegraph/api/auth.py | 20 +- .../src/pyhugegraph/api/graph.py | 8 + .../src/pyhugegraph/api/gremlin.py | 4 +- .../src/pyhugegraph/api/schema.py | 8 +- .../api/schema_manage/edge_label.py | 4 +- .../src/pyhugegraph/api/services.py | 3 +- .../src/pyhugegraph/api/traverser.py | 12 +- .../pyhugegraph/example/hugegraph_example.py | 10 +- .../src/pyhugegraph/example/hugegraph_test.py | 7 +- .../structure/vertex_label_data.py | 5 +- .../src/pyhugegraph/utils/huge_router.py | 4 +- .../src/pyhugegraph/utils/util.py | 7 +- .../src/tests/api/test_traverser.py | 32 +--- .../src/tests/client_utils.py | 32 ++-- 46 files changed, 421 insertions(+), 230 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py index 4c192c29c..9bf2ff275 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py @@ -32,9 +32,7 @@ def admin_http_api(router: APIRouter, log_stream): async def log_stream_api(req: LogStreamRequest): if admin_settings.admin_token != req.admin_token: raise generate_response( - RAGResponse( - status_code=status.HTTP_403_FORBIDDEN, # pylint: disable=E0702 - message="Invalid admin_token", + RAGResponse(status_code=status.HTTP_403_FORBIDDEN, message="Invalid admin_token") # pylint: disable=E0702 ) ) log_path = os.path.join("logs", req.log_file) diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py index f46aea02c..720c25a86 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py +++ b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py @@ -64,9 +64,13 @@ class RAGRequest(BaseModel): description="TopK results returned for each keyword \ extracted from the query, by default only the most similar one is returned.", ) +<<<<<<< HEAD client_config: Optional[GraphConfigRequest] = Query( None, description="hugegraph server config." ) +======= + client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") +>>>>>>> 87ee5d3 (style: format code with black line-length 120) # Keep prompt params in the end answer_prompt: Optional[str] = Query( @@ -163,9 +167,13 @@ class GremlinOutputType(str, Enum): class GremlinGenerateRequest(BaseModel): query: str +<<<<<<< HEAD example_num: Optional[int] = Query( 0, description="Number of Gremlin templates to use.(0 means no templates)" ) +======= + example_num: Optional[int] = Query(0, description="Number of Gremlin templates to use.(0 means no templates)") +>>>>>>> 87ee5d3 (style: format code with black line-length 120) gremlin_prompt: Optional[str] = Query( prompt.gremlin_generate_prompt, description="Prompt for the Text2Gremlin query.", diff --git a/hugegraph-llm/src/hugegraph_llm/config/index_config.py b/hugegraph-llm/src/hugegraph_llm/config/index_config.py index 21bc509cd..afe84a793 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/index_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/index_config.py @@ -33,4 +33,4 @@ class IndexConfig(BaseConfig): milvus_user: str = os.environ.get("MILVUS_USER", "") milvus_password: str = os.environ.get("MILVUS_PASSWORD", "") - now_vector_index: str = 'Faiss' + now_vector_index: str = "Faiss" diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 1b2032b23..ec01741e3 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -109,9 +109,13 @@ def create_admin_block(): ) # Error message box, initially hidden +<<<<<<< HEAD error_message = gr.Textbox( label="", visible=False, interactive=False, elem_classes="error-message" ) +======= + error_message = gr.Textbox(label="", visible=False, interactive=False, elem_classes="error-message") +>>>>>>> 87ee5d3 (style: format code with black line-length 120) # Button to submit password submit_button = gr.Button("Submit") 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 b0979763c..a8450ab5f 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -93,9 +93,7 @@ def init_rag_ui() -> gr.Interface: textbox_array_graph_config = create_configs_block() with gr.Tab(label="1. Build RAG Index 💡"): - textbox_input_text, textbox_input_schema, textbox_info_extract_template = ( - create_vector_graph_block() - ) + textbox_input_text, textbox_input_schema, textbox_info_extract_template = create_vector_graph_block() with gr.Tab(label="2. (Graph)RAG & User Functions 📖"): ( textbox_inp, @@ -104,9 +102,7 @@ def init_rag_ui() -> gr.Interface: textbox_custom_related_information, ) = create_rag_block() with gr.Tab(label="3. Text2gremlin ⚙️"): - textbox_gremlin_inp, textbox_gremlin_schema, textbox_gremlin_prompt = ( - create_text2gremlin_block() - ) + textbox_gremlin_inp, textbox_gremlin_schema, textbox_gremlin_prompt = create_text2gremlin_block() with gr.Tab(label="4. Graph Tools 🚧"): create_other_block() with gr.Tab(label="5. Admin Tools 🛠"): diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 01fc2c2da..a98c636b9 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -68,9 +68,7 @@ def test_litellm_chat(api_key, api_base, model_name, max_tokens: int) -> int: return 200 -def test_api_connection( - url, method="GET", headers=None, params=None, body=None, auth=None, origin_call=None -) -> int: +def test_api_connection(url, method="GET", headers=None, params=None, body=None, auth=None, origin_call=None) -> int: # TODO: use fastapi.request / starlette instead? log.debug("Request URL: %s", url) try: @@ -249,9 +247,13 @@ def apply_llm_config( setattr(llm_settings, f"openai_{current_llm_config}_language_model", model_name) setattr(llm_settings, f"openai_{current_llm_config}_tokens", int(max_tokens)) +<<<<<<< HEAD test_url = ( getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" ) +======= + test_url = getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" +>>>>>>> 87ee5d3 (style: format code with black line-length 120) data = { "model": model_name, "temperature": 0.01, @@ -259,6 +261,7 @@ def apply_llm_config( } <<<<<<< HEAD headers = {"Authorization": f"Bearer {api_key_or_host}"} +<<<<<<< HEAD status_code = test_api_connection( test_url, method="POST", headers=headers, body=data, origin_call=origin_call ) @@ -271,14 +274,15 @@ def apply_llm_config( arg1, arg2, arg3, settings_prefix=current_llm_config, origin_call=origin_call ) # pylint: disable=C0301 >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) +>>>>>>> 87ee5d3 (style: format code with black line-length 120) elif llm_option == "ollama/local": setattr(llm_settings, f"ollama_{current_llm_config}_host", api_key_or_host) setattr(llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port)) setattr(llm_settings, f"ollama_{current_llm_config}_language_model", model_name) - status_code = test_api_connection( - f"http://{api_key_or_host}:{api_base_or_port}", origin_call=origin_call - ) + status_code = test_api_connection(f"http://{api_key_or_host}:{api_base_or_port}", origin_call=origin_call) elif llm_option == "litellm": setattr(llm_settings, f"litellm_{current_llm_config}_api_key", api_key_or_host) @@ -286,9 +290,7 @@ def apply_llm_config( setattr(llm_settings, f"litellm_{current_llm_config}_language_model", model_name) setattr(llm_settings, f"litellm_{current_llm_config}_tokens", int(max_tokens)) - status_code = test_litellm_chat( - api_key_or_host, api_base_or_port, model_name, int(max_tokens) - ) + status_code = test_litellm_chat(api_key_or_host, api_base_or_port, model_name, int(max_tokens)) gr.Info("Configured!") llm_settings.update_env() @@ -502,16 +504,12 @@ def chat_llm_settings(llm_type): llm_config_button = gr.Button("Apply configuration") llm_config_button.click(apply_llm_config_with_chat_op, inputs=llm_config_input) # Determine whether there are Settings in the.env file - env_path = os.path.join( - os.getcwd(), ".env" - ) # Load .env from the current working directory + env_path = os.path.join(os.getcwd(), ".env") # Load .env from the current working directory env_vars = dotenv_values(env_path) api_extract_key = env_vars.get("OPENAI_EXTRACT_API_KEY") api_text2sql_key = env_vars.get("OPENAI_TEXT2GQL_API_KEY") if not api_extract_key: - llm_config_button.click( - apply_llm_config_with_text2gql_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_text2gql_op, inputs=llm_config_input) if not api_text2sql_key: <<<<<<< HEAD llm_config_button.click( @@ -524,7 +522,7 @@ def chat_llm_settings(llm_type): ======= llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) - with gr.Tab(label='mini_tasks'): + with gr.Tab(label="mini_tasks"): extract_llm_dropdown = gr.Dropdown( choices=["openai", "litellm", "qianfan_wenxin", "ollama/local"], >>>>>>> 38dce0b (feat(llm): vector db finished) @@ -663,6 +661,9 @@ def extract_llm_settings(llm_type): llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> 87ee5d3 (style: format code with black line-length 120) with gr.Tab(label="text2gql"): text2gql_llm_dropdown = gr.Dropdown( choices=["openai", "litellm", "ollama/local"], diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py index 8b78328f3..82650b907 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py @@ -58,9 +58,13 @@ def create_other_block(): async def lifespan(app: FastAPI): # pylint: disable=W0621 log.info("Starting background scheduler...") scheduler = AsyncIOScheduler() +<<<<<<< HEAD scheduler.add_job( backup_data, trigger=CronTrigger(hour=1, minute=0), id="daily_backup", replace_existing=True ) +======= + scheduler.add_job(backup_data, trigger=CronTrigger(hour=1, minute=0), id="daily_backup", replace_existing=True) +>>>>>>> 87ee5d3 (style: format code with black line-length 120) scheduler.start() log.info("Starting vid embedding update task...") diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index f61b57123..cf40c444e 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -383,7 +383,9 @@ def toggle_slider(enable): 0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False ) - graph_vector_radio.change(toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio) # pylint: disable=no-member + graph_vector_radio.change( + toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio + ) # pylint: disable=no-member near_neighbor_first = gr.Checkbox( value=False, label="Near neighbor first(Optional)", diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 7ebed651a..32529fb1e 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -60,12 +60,16 @@ def error(cls, message: str) -> "GremlinResult": @classmethod def success_result( +<<<<<<< HEAD cls, match_result: str, template_gremlin: str, raw_gremlin: str, template_exec: str, raw_exec: str, +======= + cls, match_result: str, template_gremlin: str, raw_gremlin: str, template_exec: str, raw_exec: str +>>>>>>> 87ee5d3 (style: format code with black line-length 120) ) -> "GremlinResult": """Create a successful result""" return cls( @@ -101,8 +105,12 @@ def build_example_vector_index(temp_file) -> dict: os.makedirs(index_path) ======= vector_index = get_vector_index_class(index_settings.now_vector_index) +<<<<<<< HEAD assert vector_index, 'vector db name is error' >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + assert vector_index, "vector db name is error" +>>>>>>> 87ee5d3 (style: format code with black line-length 120) if temp_file is None: full_path = os.path.join(resource_path, "demo", "text2gremlin.csv") else: @@ -115,6 +123,10 @@ def build_example_vector_index(temp_file) -> dict: try: import shutil +<<<<<<< HEAD +======= + os.makedirs(os.path.dirname(target_file), exist_ok=True) +>>>>>>> 87ee5d3 (style: format code with black line-length 120) shutil.copy2(full_path, target_file) log.info("Successfully copied file to: %s", target_file) except (OSError, IOError) as e: @@ -296,11 +308,19 @@ def gremlin_generate_for_ui(inp, example_num, schema, gremlin_prompt): ) return ( +<<<<<<< HEAD match_result_str, res.get("template_gremlin", "") or "", res.get("raw_gremlin", "") or "", res.get("template_execution_result", "") or "", res.get("raw_execution_result", "") or "", +======= + result.match_result, + result.template_gremlin or "", + result.raw_gremlin or "", + result.template_exec_result or "", + result.raw_exec_result or "", +>>>>>>> 87ee5d3 (style: format code with black line-length 120) ) @@ -435,9 +455,7 @@ def gremlin_generate_selective( if not requested_outputs: # None or empty list requested_outputs = output_keys - result = gremlin_generate( - inp, example_num, schema_input, gremlin_prompt_input, requested_outputs - ) + result = gremlin_generate(inp, example_num, schema_input, gremlin_prompt_input, requested_outputs) outputs_dict: Dict[str, Any] = {} diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 8b081cf9a..0efc11d59 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -58,11 +58,7 @@ def store_prompt(doc, schema, example_prompt): # update env variables: doc, schema and example_prompt - if ( - prompt.doc_input_text != doc - or prompt.graph_schema != schema - or prompt.extract_graph_prompt != example_prompt - ): + if prompt.doc_input_text != doc or prompt.graph_schema != schema or prompt.extract_graph_prompt != example_prompt: prompt.doc_input_text = doc prompt.graph_schema = schema prompt.extract_graph_prompt = example_prompt @@ -104,11 +100,15 @@ def load_query_examples(): language = getattr( prompt, "language", +<<<<<<< HEAD ( getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") else "EN" ), +======= + getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") else "EN", +>>>>>>> 87ee5d3 (style: format code with black line-length 120) ) if language.upper() == "CN": examples_path = os.path.join(resource_path, "prompt_examples", "query_examples_CN.json") @@ -509,3 +509,165 @@ async def timely_update_vid_embedding(interval_seconds: int = 3600): except Exception as e: log.warning("Failed to execute update_vid_embedding: %s", e, exc_info=True) await asyncio.sleep(interval_seconds) +<<<<<<< HEAD +======= + + +def create_vector_graph_block(): + # pylint: disable=no-member + # pylint: disable=C0301 + # pylint: disable=unexpected-keyword-arg + gr.Markdown( + """## Build Vector/Graph Index & Extract Knowledge Graph +- Docs: + - text: Build rag index from plain text + - file: Upload file(s) which should be TXT or .docx (Multiple files can be selected together) +- [Schema](https://hugegraph.apache.org/docs/clients/restful-api/schema/): (Accept **2 types**) + - User-defined Schema (JSON format, follow the [template](https://github.com/apache/incubator-hugegraph-ai/blob/aff3bbe25fa91c3414947a196131be812c20ef11/hugegraph-llm/src/hugegraph_llm/config/config_data.py#L125) + to modify it) + - Specify the name of the HugeGraph graph instance, it will automatically get the schema from it (like + **"hugegraph"**) +- Graph Extract Prompt Header: The user-defined prompt of graph extracting +- If already exist the graph data, you should click "**Rebuild vid Index**" to update the index +""" + ) + + with gr.Row(): + with gr.Column(): + with gr.Tab("text") as tab_upload_text: + input_text = gr.Textbox( + value=prompt.doc_input_text, label="Input Doc(s)", lines=20, show_copy_button=True + ) + with gr.Tab("file") as tab_upload_file: + input_file = gr.File( + value=None, + label="Docs (multi-files can be selected together)", + file_count="multiple", + ) + input_schema = gr.Code(value=prompt.graph_schema, label="Graph Schema", language="json", lines=15, max_lines=29) + info_extract_template = gr.Code( + value=prompt.extract_graph_prompt, + label="Graph Extract Prompt Header", + language="markdown", + lines=15, + max_lines=29, + ) + + out = gr.Code(label="Output Info", language="json", elem_classes="code-container-edit") + + with gr.Row(): + with gr.Accordion("Get RAG Info", open=False): + with gr.Column(): + vector_index_btn0 = gr.Button("Get Vector Index Info", size="sm") + graph_index_btn0 = gr.Button("Get Graph Index Info", size="sm") + with gr.Accordion("Clear RAG Data", open=False): + with gr.Column(): + vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") + graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") + graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") + vector_import_bt = gr.Button("Import into Vector", variant="primary") + graph_extract_bt = gr.Button("Extract Graph Data (1)", variant="primary") + graph_loading_bt = gr.Button("Load into GraphDB (2)", interactive=True) + graph_index_rebuild_bt = gr.Button("Update Vid Embedding") + + gr.Markdown("---") + with gr.Accordion("Graph Schema Generator", open=False): + gr.Markdown( + "Provide **query examples** and **few-shot examples**, " + "then click **Generate Schema** to automatically create graph schema." + ) + with gr.Row(): + query_example = gr.Code( + value=load_query_examples(), + label="Query Examples", + language="json", + lines=10, + max_lines=15, + ) + few_shot = gr.Code( + value=load_schema_fewshot_examples(), + label="Few-shot Example", + language="json", + lines=10, + max_lines=15, + ) + build_schema_bt = gr.Button("Generate Schema", variant="primary") + + # 事件绑定 + vector_index_btn0.click(get_vector_index_info, outputs=out).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) + vector_index_btn1.click(clean_vector_index).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) + vector_import_bt.click(build_vector_index, inputs=[input_file, input_text], outputs=out).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) + graph_index_btn0.click(get_graph_index_info, outputs=out).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) + graph_index_btn1.click(clean_all_graph_index).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) + graph_data_btn0.click(clean_all_graph_data).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) + graph_index_rebuild_bt.click(update_vid_embedding, outputs=out).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) + + graph_extract_bt.click( + extract_graph, inputs=[input_file, input_text, input_schema, info_extract_template], outputs=[out] + ).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) + + graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( + update_vid_embedding + ).then( + store_prompt, + inputs=[input_text, input_schema, info_extract_template], + ) + + build_schema_bt.click( + lambda it, qe, fs: extract_graph([], it, prompt.graph_schema, prompt.extract_graph_prompt), + inputs=[input_text, query_example, few_shot], + outputs=[input_schema], + ).then( + store_prompt, + inputs=[ + input_text, + input_schema, + info_extract_template, + ], + ) + + def on_tab_select(input_f, input_t, evt: gr.SelectData): + print(f"You selected {evt.value} at {evt.index} from {evt.target}") + if evt.value == "file": + return input_f, "" + if evt.value == "text": + return [], input_t + return [], "" + + tab_upload_file.select( + fn=on_tab_select, + inputs=[input_file, input_text], + outputs=[input_file, input_text], + ) + tab_upload_text.select( + fn=on_tab_select, + inputs=[input_file, input_text], + outputs=[input_file, input_text], + ) + + return input_text, input_schema, info_extract_template +>>>>>>> 87ee5d3 (style: format code with black line-length 120) diff --git a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py index ee173b284..12572a9e7 100644 --- a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py +++ b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py @@ -22,9 +22,13 @@ class ChunkSplitter: def __init__( +<<<<<<< HEAD self, split_type: Literal["paragraph", "sentence"] = "paragraph", language: Literal["zh", "en"] = "zh", +======= + self, split_type: Literal["paragraph", "sentence"] = "paragraph", language: Literal["zh", "en"] = "zh" +>>>>>>> 87ee5d3 (style: format code with black line-length 120) ): if language == "zh": separators = ["\n\n", "\n", "。", "，", ""] @@ -33,6 +37,7 @@ def __init__( else: raise ValueError("Argument `language` must be zh or en!") if split_type == "paragraph": +<<<<<<< HEAD self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=500, chunk_overlap=30, separators=separators ) @@ -40,6 +45,11 @@ def __init__( self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=50, chunk_overlap=0, separators=separators ) +======= + self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=30, separators=separators) + elif split_type == "sentence": + self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=50, chunk_overlap=0, separators=separators) +>>>>>>> 87ee5d3 (style: format code with black line-length 120) else: raise ValueError("Arg `type` must be paragraph, sentence!") diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py index acdbdf511..631e7908b 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -224,7 +224,7 @@ def get_vector_index_info(self) -> dict: @staticmethod def clean(*name: str): - name_str = '_'.join(name) + name_str = "_".join(name) connections.connect( host=index_settings.milvus_host, port=index_settings.milvus_port, @@ -236,7 +236,7 @@ def clean(*name: str): @staticmethod def from_name(embed_dim: int, *name: str) -> "MilvusVectorIndex": - name_str = '_'.join(name) + name_str = "_".join(name) assert index_settings.milvus_host, "Qdrant host is not configured" return MilvusVectorIndex( name_str, @@ -249,7 +249,7 @@ def from_name(embed_dim: int, *name: str) -> "MilvusVectorIndex": @staticmethod def exist(*name: str) -> bool: - name_str = '_'.join(name) + name_str = "_".join(name) connections.connect( host=index_settings.milvus_host, port=index_settings.milvus_port, diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index 98f2eaf9c..7e403f312 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -179,7 +179,7 @@ def get_vector_index_info(self) -> Dict: @staticmethod def clean(*name: str): - name_str = '_'.join(name) + name_str = "_".join(name) client = QdrantClient( host=index_settings.qdrant_host, port=index_settings.qdrant_port, api_key=index_settings.qdrant_api_key ) @@ -192,7 +192,7 @@ def clean(*name: str): @staticmethod def from_name(embed_dim: int, *name: str) -> "QdrantVectorIndex": assert index_settings.qdrant_host, "Qdrant host is not configured" - name_str = '_'.join(name) + name_str = "_".join(name) return QdrantVectorIndex( name=name_str, host=index_settings.qdrant_host, @@ -203,7 +203,7 @@ def from_name(embed_dim: int, *name: str) -> "QdrantVectorIndex": @staticmethod def exist(*name: str) -> bool: - name_str = '_'.join(name) + name_str = "_".join(name) client = QdrantClient( host=index_settings.qdrant_host, port=index_settings.qdrant_port, api_key=index_settings.qdrant_api_key ) diff --git a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py index 7e93d0ec0..f13e11e6f 100644 --- a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py +++ b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py @@ -46,12 +46,7 @@ async def dispatch(self, request: Request, call_next): "%s - Args: %s, IP: %s, URL: %s", request.method, request.query_params, -<<<<<<< HEAD - request.client.host, + request.client.host, # type: ignore request.url, -======= - request.client.host, # type: ignore - request.url ->>>>>>> 38dce0b (feat(llm): vector db finished) ) return response diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index 8a96d7182..d96840911 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -65,7 +65,7 @@ def __init__(self): def get_embedding(self): if self.embedding_type == "openai": - assert llm_settings.openai_embedding_model_dim, 'openai_embedding_model_dim is need' + assert llm_settings.openai_embedding_model_dim, "openai_embedding_model_dim is need" return OpenAIEmbedding( embedding_dimension=llm_settings.openai_embedding_model_dim, model_name=llm_settings.openai_embedding_model, @@ -73,7 +73,7 @@ def get_embedding(self): api_base=llm_settings.openai_embedding_api_base, ) if self.embedding_type == "ollama/local": - assert llm_settings.ollama_embedding_model_dim, 'ollama_embedding_model_dim is need' + assert llm_settings.ollama_embedding_model_dim, "ollama_embedding_model_dim is need" return OllamaEmbedding( <<<<<<< HEAD model_name=llm_settings.ollama_embedding_model, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py index da99d1e65..a29171ec5 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py @@ -29,7 +29,7 @@ def __init__(self, model_name: str, host: str = "127.0.0.1", port: int = 11434, ======= def __init__( self, - model: str = 'quentinz/bge-large-zh-v1.5', + model: str = "quentinz/bge-large-zh-v1.5", embedding_dimension: int = 1024, host: str = "127.0.0.1", port: int = 11434, diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py index 6d08ce8cd..1ca304baa 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py @@ -118,9 +118,13 @@ async def agenerate_streaming( messages = [{"role": "user", "content": prompt}] try: +<<<<<<< HEAD async_generator = await self.async_client.chat( model=self.model, messages=messages, stream=True ) +======= + async_generator = await self.async_client.chat(model=self.model, messages=messages, stream=True) +>>>>>>> 87ee5d3 (style: format code with black line-length 120) async for chunk in async_generator: token = chunk.get("message", {}).get("content", "") if on_token_callback: 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 c31e77af7..a8530632b 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 @@ -56,9 +56,7 @@ def _get_text_splitter(self, split_type: str): chunk_size=500, chunk_overlap=30, separators=self.separators ).split_text if split_type == SPLIT_TYPE_SENTENCE: - return RecursiveCharacterTextSplitter( - chunk_size=50, chunk_overlap=0, separators=self.separators - ).split_text + return RecursiveCharacterTextSplitter(chunk_size=50, chunk_overlap=0, separators=self.separators).split_text raise ValueError("Type must be paragraph, sentence, html or markdown") def run(self, context: Optional[Dict[str, Any]]) -> Dict[str, Any]: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py index 0d9967020..a005c7472 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py @@ -77,8 +77,6 @@ def _filter_keywords( results.add(token) sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: - results.update( - {w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language)} - ) + results.update({w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language)}) return list(results) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index b37707ca3..f972c79b8 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -250,6 +250,7 @@ def run(self, **kwargs) -> Dict[str, Any]: :return: Final context after all operators have been executed. """ if len(self._operators) == 0: +<<<<<<< HEAD <<<<<<< HEAD self.extract_keywords().query_graphdb( max_graph_items=kwargs.get("max_graph_items") @@ -257,6 +258,9 @@ def run(self, **kwargs) -> Dict[str, Any]: ======= self.extract_keywords().query_graphdb(max_graph_items=kwargs.get('max_graph_items')).synthesize_answer() >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + self.extract_keywords().query_graphdb(max_graph_items=kwargs.get("max_graph_items")).synthesize_answer() +>>>>>>> 87ee5d3 (style: format code with black line-length 120) context = kwargs diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py index 9138f9e9b..df697aad0 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py @@ -62,9 +62,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = ( - f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") - ) + context_str = f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -74,13 +72,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: graph_result_context, vector_result_context = self.handle_vector_graph(context) context = asyncio.run( - self.async_generate( - context, - context_head_str, - context_tail_str, - vector_result_context, - graph_result_context, - ) + self.async_generate(context, context_head_str, context_tail_str, vector_result_context, graph_result_context) ) return context @@ -119,9 +111,7 @@ async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[st context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = ( - f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") - ) + context_str = f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -151,11 +141,7 @@ async def async_generate( final_prompt = self._question async_tasks["raw_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) if self._vector_only_answer: - context_str = ( - f"{context_head_str}\n" - f"{vector_result_context}\n" - f"{context_tail_str}".strip("\n") - ) + context_str = f"{context_head_str}\n" f"{vector_result_context}\n" f"{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -164,11 +150,7 @@ async def async_generate( self._llm.agenerate(prompt=final_prompt) ) if self._graph_only_answer: - context_str = ( - f"{context_head_str}\n" - f"{graph_result_context}\n" - f"{context_tail_str}".strip("\n") - ) + context_str = f"{context_head_str}\n" f"{graph_result_context}\n" f"{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -180,16 +162,10 @@ async def async_generate( context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = ( - f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") - ) + context_str = f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") - final_prompt = self._prompt_template.format( - context_str=context_str, query_str=self._question - ) - async_tasks["graph_vector_task"] = asyncio.create_task( - self._llm.agenerate(prompt=final_prompt) - ) + final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + async_tasks["graph_vector_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) async_tasks_mapping = { "raw_task": "raw_answer", @@ -204,14 +180,7 @@ async def async_generate( context[context_key] = response log.debug("Query Answer: %s", response) - ops = sum( - [ - self._raw_answer, - self._vector_only_answer, - self._graph_only_answer, - self._graph_vector_answer, - ] - ) + ops = sum([self._raw_answer, self._vector_only_answer, self._graph_only_answer, self._graph_vector_answer]) context["call_count"] = context.get("call_count", 0) + ops return context @@ -235,11 +204,7 @@ async def async_streaming_generate( ) auto_id += 1 if self._vector_only_answer: - context_str = ( - f"{context_head_str}\n" - f"{vector_result_context}\n" - f"{context_tail_str}".strip("\n") - ) + context_str = f"{context_head_str}\n" f"{vector_result_context}\n" f"{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -251,11 +216,7 @@ async def async_streaming_generate( ) auto_id += 1 if self._graph_only_answer: - context_str = ( - f"{context_head_str}\n" - f"{graph_result_context}\n" - f"{context_tail_str}".strip("\n") - ) + context_str = f"{context_head_str}\n" f"{graph_result_context}\n" f"{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -270,9 +231,7 @@ async def async_streaming_generate( context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = ( - f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") - ) + context_str = f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -284,14 +243,7 @@ async def async_streaming_generate( ) auto_id += 1 - ops = sum( - [ - self._raw_answer, - self._vector_only_answer, - self._graph_only_answer, - self._graph_vector_answer, - ] - ) + ops = sum([self._raw_answer, self._vector_only_answer, self._graph_only_answer, self._graph_vector_answer]) context["call_count"] = context.get("call_count", 0) + ops async_tasks = [asyncio.create_task(anext(gen)) for gen in async_generators] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py index 2ac2eafff..981123954 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py @@ -51,10 +51,14 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: llm_output = self.llm.generate(prompt=prompt) data["triples"] = [] extract_triples_by_regex(llm_output, data) +<<<<<<< HEAD print( f"LLM {self.__class__.__name__} input:{prompt} \n" f" output: {llm_output} \n data: {data}" ) +======= + print(f"LLM {self.__class__.__name__} input:{prompt} \n" f" output: {llm_output} \n data: {data}") +>>>>>>> 87ee5d3 (style: format code with black line-length 120) data["call_count"] = data.get("call_count", 0) + 1 return data 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 8897e0fea..504d683df 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 @@ -76,8 +76,12 @@ def generate_extract_triple_prompt(text, schema=None) -> str: if schema: return schema_real_prompt log.warning( +<<<<<<< HEAD "Recommend to provide a graph schema to improve the extraction accuracy. " "Now using the default schema." +======= + "Recommend to provide a graph schema to improve the extraction accuracy. " "Now using the default schema." +>>>>>>> 87ee5d3 (style: format code with black line-length 120) ) return text_based_prompt @@ -107,9 +111,13 @@ def extract_triples_by_regex_with_schema(schema, text, graph): # TODO: use a more efficient way to compare the extract & input property p_lower = p.lower() for vertex in schema["vertices"]: +<<<<<<< HEAD if vertex["vertex_label"] == label and any( pp.lower() == p_lower for pp in vertex["properties"] ): +======= + if vertex["vertex_label"] == label and any(pp.lower() == p_lower for pp in vertex["properties"]): +>>>>>>> 87ee5d3 (style: format code with black line-length 120) id = f"{label}-{s}" if id not in vertices_dict: vertices_dict[id] = { @@ -126,6 +134,7 @@ def extract_triples_by_regex_with_schema(schema, text, graph): source_label = edge["source_vertex_label"] source_id = f"{source_label}-{s}" if source_id not in vertices_dict: +<<<<<<< HEAD vertices_dict[source_id] = { "id": source_id, "name": s, @@ -149,6 +158,14 @@ def extract_triples_by_regex_with_schema(schema, text, graph): "properties": {}, } ) +======= + vertices_dict[source_id] = {"id": source_id, "name": s, "label": source_label, "properties": {}} + target_label = edge["target_vertex_label"] + target_id = f"{target_label}-{o}" + if target_id not in vertices_dict: + vertices_dict[target_id] = {"id": target_id, "name": o, "label": target_label, "properties": {}} + graph["edges"].append({"start": source_id, "end": target_id, "type": label, "properties": {}}) +>>>>>>> 87ee5d3 (style: format code with black line-length 120) break graph["vertices"] = list(vertices_dict.values()) @@ -199,7 +216,11 @@ def valid(self, element_id: str, max_length: int = 256) -> bool: def _filter_long_id(self, graph) -> Dict[str, List[Any]]: graph["vertices"] = [vertex for vertex in graph["vertices"] if self.valid(vertex["id"])] +<<<<<<< HEAD graph["edges"] = [ edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"]) ] +======= + graph["edges"] = [edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"])] +>>>>>>> 87ee5d3 (style: format code with black line-length 120) return graph diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index 425f2a70b..271b530c6 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -64,9 +64,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: end_time = time.perf_counter() log.debug("Keyword extraction time: %.2f seconds", end_time - start_time) - keywords = self._extract_keywords_from_response( - response=response, lowercase=False, start_token="KEYWORDS:" - ) + keywords = self._extract_keywords_from_response(response=response, lowercase=False, start_token="KEYWORDS:") keywords = {k.replace("'", "") for k in keywords} context["keywords"] = list(keywords) log.info("User Query: %s\nKeywords: %s", self._query, context["keywords"]) @@ -87,11 +85,15 @@ def _extract_keywords_from_response( for match in matches: match = match[len(start_token) :].strip() +<<<<<<< HEAD keywords.extend( k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1 ) +======= + keywords.extend(k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1) +>>>>>>> 87ee5d3 (style: format code with black line-length 120) # if the keyword consists of multiple words, split into sub-words (removing stopwords) results = set(keywords) 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 565d79023..0b34717e3 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 @@ -125,8 +125,12 @@ 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( +<<<<<<< HEAD "Invalid property graph! No JSON object found, " "please check the output format example in prompt." +======= + "Invalid property graph! No JSON object found, " "please check the output format example in prompt." +>>>>>>> 87ee5d3 (style: format code with black line-length 120) ) return [] json_str = json_match.group(1).strip() diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py index 928948413..1e33514ca 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py @@ -86,10 +86,14 @@ def _extract_schema(self, response: str) -> Dict[str, Any]: raise RuntimeError("Invalid JSON response from LLM") from e def build_prompt( +<<<<<<< HEAD self, raw_texts: List[str], query_examples: List[Dict[str, str]], few_shot_schema: Dict[str, Any], +======= + self, raw_texts: List[str], query_examples: List[Dict[str, str]], few_shot_schema: Dict[str, Any] +>>>>>>> 87ee5d3 (style: format code with black line-length 120) ) -> str: return self.schema_prompt.format( raw_texts=self._format_raw_texts(raw_texts), diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py index 38eabb16e..6beeb0291 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py @@ -20,10 +20,7 @@ import re REGEX = ( - r"Nodes:\s+(.*?)\s?\s?" - r"Relationships:\s?\s?" - r"NodesSchemas:\s+(.*?)\s?\s?" - r"RelationshipsSchemas:\s?\s?(.*)" + r"Nodes:\s+(.*?)\s?\s?" r"Relationships:\s?\s?" r"NodesSchemas:\s+(.*?)\s?\s?" r"RelationshipsSchemas:\s?\s?(.*)" ) INTERNAL_REGEX = r"\[(.*?)\]" JSON_REGEX = r"\{.*\}" diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 70c94b7bf..ddb531508 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -66,20 +66,33 @@ def get_graph_index_info(): vector_index_info = vector_index_entity.get_vector_index_info() >>>>>>> 38dce0b (feat(llm): vector db finished) graph_summary_info["vid_index"] = { - "embed_dim": vector_index_info['embed_dim'], - "num_vectors": vector_index_info['vector_info']['chunk_vector_num'], - "num_vids": vector_index_info['vector_info']['graph_properties_vector_num'], + "embed_dim": vector_index_info["embed_dim"], + "num_vectors": vector_index_info["vector_info"]["chunk_vector_num"], + "num_vids": vector_index_info["vector_info"]["graph_properties_vector_num"], } return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) def clean_all_graph_index(): +<<<<<<< HEAD folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None) ) FaissVectorIndex.clean(str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix) FaissVectorIndex.clean(str(os.path.join(resource_path, folder_name, "gremlin_examples")), filename_prefix) +======= + # 清理 Faiss 索引目录（兼容默认） + faiss_chunks = os.path.join(resource_path, huge_settings.graph_name, "graph_vids") + if os.path.isdir(faiss_chunks): + from ..indices.vector_index.faiss_vector_store import FaissVectorIndex + + FaissVectorIndex.clean(huge_settings.graph_name, "graph_vids") + # 清理默认的 gremlin_examples + from ..indices.vector_index.faiss_vector_store import FaissVectorIndex + + FaissVectorIndex.clean("gremlin_examples") +>>>>>>> 87ee5d3 (style: format code with black line-length 120) log.warning("Clear graph index and text2gql index successfully!") gr.Info("Clear graph index and text2gql index successfully!") @@ -187,7 +200,31 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: def build_schema(input_text, query_example, few_shot): +<<<<<<< HEAD scheduler = SchedulerSingleton.get_instance() +======= + context = {"raw_texts": [input_text] if input_text else [], "query_examples": [], "few_shot_schema": {}} + + if few_shot: + try: + context["few_shot_schema"] = json.loads(few_shot) + except json.JSONDecodeError as e: + raise gr.Error(f"Few Shot Schema is not in a valid JSON format: {e}") from e + + if query_example: + try: + parsed_examples = json.loads(query_example) + # Validate and retain the description and gremlin fields + context["query_examples"] = [ + {"description": ex.get("description", ""), "gremlin": ex.get("gremlin", "")} + for ex in parsed_examples + if isinstance(ex, dict) and "description" in ex and "gremlin" in ex + ] + except json.JSONDecodeError as e: + raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e + + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) +>>>>>>> 87ee5d3 (style: format code with black line-length 120) try: return scheduler.schedule_flow( "build_schema", input_text, query_example, few_shot diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index 147c0074c..28b738b13 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -140,10 +140,14 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): elif filename == "vertices.json": data_full = client.gremlin().exec(query)["data"][0]["vertices"] data = ( +<<<<<<< HEAD [ {key: value for key, value in vertex.items() if key != "id"} for vertex in data_full ] +======= + [{key: value for key, value in vertex.items() if key != "id"} for vertex in data_full] +>>>>>>> 87ee5d3 (style: format code with black line-length 120) if all_pk_flag else data_full ) @@ -164,9 +168,7 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): def manage_backup_retention(): try: backup_dirs = [ - os.path.join(BACKUP_DIR, d) - for d in os.listdir(BACKUP_DIR) - if os.path.isdir(os.path.join(BACKUP_DIR, d)) + os.path.join(BACKUP_DIR, d) for d in os.listdir(BACKUP_DIR) if os.path.isdir(os.path.join(BACKUP_DIR, d)) ] backup_dirs.sort(key=os.path.getctime) if len(backup_dirs) > MAX_BACKUP_DIRS: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/log.py b/hugegraph-llm/src/hugegraph_llm/utils/log.py index b64017454..aef36e449 100755 --- a/hugegraph-llm/src/hugegraph_llm/utils/log.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/log.py @@ -27,11 +27,15 @@ # Initialize the root logger first with Rich handler root_logger = init_logger( +<<<<<<< HEAD log_output=LOG_FILE, log_level=INFO, logger_name="root", propagate_logs=True, stdout_logging=True, +======= + log_output=LOG_FILE, log_level=INFO, logger_name="root", propagate_logs=True, stdout_logging=True +>>>>>>> 87ee5d3 (style: format code with black line-length 120) ) # Initialize custom logger diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index d8c87fdd6..4a1363676 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -78,7 +78,7 @@ def get_vector_index_info(): ) return json.dumps( - {**vector_index_entity.get_vector_index_info(), 'now_vector_index': index_settings.now_vector_index}, + {**vector_index_entity.get_vector_index_info(), "now_vector_index": index_settings.now_vector_index}, ensure_ascii=False, indent=2, ) diff --git a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py index 7ad914468..76fd4ccd1 100644 --- a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py +++ b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py @@ -32,6 +32,10 @@ def test_stream_generate(self): def on_token_callback(chunk): print(chunk, end="", flush=True) +<<<<<<< HEAD ollama_client.generate_streaming( prompt="What is the capital of France?", on_token_callback=on_token_callback ) +======= + ollama_client.generate_streaming(prompt="What is the capital of France?", on_token_callback=on_token_callback) +>>>>>>> 87ee5d3 (style: format code with black line-length 120) diff --git a/hugegraph-python-client/src/pyhugegraph/api/auth.py b/hugegraph-python-client/src/pyhugegraph/api/auth.py index ab7d66169..ea3695b99 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/auth.py +++ b/hugegraph-python-client/src/pyhugegraph/api/auth.py @@ -31,9 +31,7 @@ def list_users(self, limit=None): return self._invoke_request(params=params) @router.http("POST", "auth/users") - def create_user( - self, user_name, user_password, user_phone=None, user_email=None - ) -> Optional[Dict]: + def create_user(self, user_name, user_password, user_phone=None, user_email=None) -> Optional[Dict]: return self._invoke_request( data=json.dumps( { @@ -118,9 +116,7 @@ def revoke_accesses(self, access_id) -> Optional[Dict]: # pylint: disable=unuse return self._invoke_request() @router.http("PUT", "auth/accesses/{access_id}") - def modify_accesses( - self, access_id, access_description # pylint: disable=unused-argument - ) -> Optional[Dict]: + def modify_accesses(self, access_id, access_description) -> Optional[Dict]: # pylint: disable=unused-argument # The permission of access can\'t be updated data = {"access_description": access_description} return self._invoke_request(data=json.dumps(data)) @@ -134,9 +130,7 @@ def list_accesses(self) -> Optional[Dict]: return self._invoke_request() @router.http("POST", "auth/targets") - def create_target( - self, target_name, target_graph, target_url, target_resources - ) -> Optional[Dict]: + def create_target(self, target_name, target_graph, target_url, target_resources) -> Optional[Dict]: return self._invoke_request( data=json.dumps( { @@ -173,9 +167,7 @@ def update_target( ) @router.http("GET", "auth/targets/{target_id}") - def get_target( - self, target_id, response=None # pylint: disable=unused-argument - ) -> Optional[Dict]: + def get_target(self, target_id, response=None) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "auth/targets") @@ -192,9 +184,7 @@ def delete_belong(self, belong_id) -> None: # pylint: disable=unused-argument return self._invoke_request() @router.http("PUT", "auth/belongs/{belong_id}") - def update_belong( - self, belong_id, description # pylint: disable=unused-argument - ) -> Optional[Dict]: + def update_belong(self, belong_id, description) -> Optional[Dict]: # pylint: disable=unused-argument data = {"belong_description": description} return self._invoke_request(data=json.dumps(data)) diff --git a/hugegraph-python-client/src/pyhugegraph/api/graph.py b/hugegraph-python-client/src/pyhugegraph/api/graph.py index 4555eeda4..94593836a 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/graph.py +++ b/hugegraph-python-client/src/pyhugegraph/api/graph.py @@ -138,17 +138,25 @@ def addEdges(self, input_data) -> Optional[List[EdgeData]]: return None @router.http("PUT", "graph/edges/{edge_id}?action=append") +<<<<<<< HEAD def appendEdge( self, edge_id, properties # pylint: disable=unused-argument ) -> Optional[EdgeData]: +======= + def appendEdge(self, edge_id, properties) -> Optional[EdgeData]: # pylint: disable=unused-argument +>>>>>>> 87ee5d3 (style: format code with black line-length 120) if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None @router.http("PUT", "graph/edges/{edge_id}?action=eliminate") +<<<<<<< HEAD def eliminateEdge( self, edge_id, properties # pylint: disable=unused-argument ) -> Optional[EdgeData]: +======= + def eliminateEdge(self, edge_id, properties) -> Optional[EdgeData]: # pylint: disable=unused-argument +>>>>>>> 87ee5d3 (style: format code with black line-length 120) if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None diff --git a/hugegraph-python-client/src/pyhugegraph/api/gremlin.py b/hugegraph-python-client/src/pyhugegraph/api/gremlin.py index e02e7fb2a..d91bdb04b 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/gremlin.py +++ b/hugegraph-python-client/src/pyhugegraph/api/gremlin.py @@ -43,9 +43,7 @@ def exec(self, gremlin): try: if response := self._invoke_request(data=gremlin_data.to_json()): return ResponseData(response).result - log.error( # pylint: disable=logging-fstring-interpolation - f"Gremlin can't get results: {str(response)}" - ) + log.error(f"Gremlin can't get results: {str(response)}") # pylint: disable=logging-fstring-interpolation return None except Exception as e: raise NotFoundError(f"Gremlin can't get results: {e}") from e diff --git a/hugegraph-python-client/src/pyhugegraph/api/schema.py b/hugegraph-python-client/src/pyhugegraph/api/schema.py index 7e8926678..3576ce66b 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/schema.py +++ b/hugegraph-python-client/src/pyhugegraph/api/schema.py @@ -68,9 +68,7 @@ def getSchema(self, _format: str = "json") -> Optional[Dict]: # pylint: disable return self._invoke_request() @router.http("GET", "schema/propertykeys/{property_name}") - def getPropertyKey( - self, property_name # pylint: disable=unused-argument - ) -> Optional[PropertyKeyData]: + def getPropertyKey(self, property_name) -> Optional[PropertyKeyData]: # pylint: disable=unused-argument if response := self._invoke_request(): return PropertyKeyData(response) return None @@ -95,9 +93,7 @@ def getVertexLabels(self) -> Optional[List[VertexLabelData]]: return None @router.http("GET", "schema/edgelabels/{label_name}") - def getEdgeLabel( - self, label_name: str # pylint: disable=unused-argument - ) -> Optional[EdgeLabelData]: + def getEdgeLabel(self, label_name: str) -> Optional[EdgeLabelData]: # pylint: disable=unused-argument if response := self._invoke_request(): return EdgeLabelData(response) log.error("EdgeLabel not found: %s", str(response)) diff --git a/hugegraph-python-client/src/pyhugegraph/api/schema_manage/edge_label.py b/hugegraph-python-client/src/pyhugegraph/api/schema_manage/edge_label.py index 93f218001..91608fe6a 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/schema_manage/edge_label.py +++ b/hugegraph-python-client/src/pyhugegraph/api/schema_manage/edge_label.py @@ -150,9 +150,7 @@ def append(self): def eliminate(self): name = self._parameter_holder.get_value("name") user_data = ( - self._parameter_holder.get_value("user_data") - if self._parameter_holder.get_value("user_data") - else {} + self._parameter_holder.get_value("user_data") if self._parameter_holder.get_value("user_data") else {} ) path = f"schema/edgelabels/{name}?action=eliminate" data = {"name": name, "user_data": user_data} diff --git a/hugegraph-python-client/src/pyhugegraph/api/services.py b/hugegraph-python-client/src/pyhugegraph/api/services.py index f353673db..4fac4aa69 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/services.py +++ b/hugegraph-python-client/src/pyhugegraph/api/services.py @@ -125,7 +125,6 @@ def delete_service(self, graphspace: str, service: str): # pylint: disable=unus None """ return self._sess.request( - f"/graphspaces/{graphspace}/services/{service}" - f"?confirm_message=I'm sure to delete the service", + f"/graphspaces/{graphspace}/services/{service}" f"?confirm_message=I'm sure to delete the service", "DELETE", ) diff --git a/hugegraph-python-client/src/pyhugegraph/api/traverser.py b/hugegraph-python-client/src/pyhugegraph/api/traverser.py index 72dddb07a..199ed0167 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/traverser.py +++ b/hugegraph-python-client/src/pyhugegraph/api/traverser.py @@ -49,9 +49,7 @@ def shortest_path(self, source_id, target_id, max_depth): # pylint: disable=unu "GET", 'traversers/allshortestpaths?source="{source_id}"&target="{target_id}"&max_depth={max_depth}', ) - def all_shortest_paths( - self, source_id, target_id, max_depth # pylint: disable=unused-argument - ): + def all_shortest_paths(self, source_id, target_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() @router.http( @@ -59,9 +57,7 @@ def all_shortest_paths( 'traversers/weightedshortestpath?source="{source_id}"&target="{target_id}"' "&weight={weight}&max_depth={max_depth}", ) - def weighted_shortest_path( - self, source_id, target_id, weight, max_depth # pylint: disable=unused-argument - ): + def weighted_shortest_path(self, source_id, target_id, weight, max_depth): # pylint: disable=unused-argument return self._invoke_request() @router.http( @@ -130,9 +126,7 @@ def advanced_paths( ) @router.http("POST", "traversers/customizedpaths") - def customized_paths( - self, sources, steps, sort_by="INCR", with_vertex=True, capacity=-1, limit=-1 - ): + def customized_paths(self, sources, steps, sort_by="INCR", with_vertex=True, capacity=-1, limit=-1): return self._invoke_request( data=json.dumps( { diff --git a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py index d5cc0eb9d..5026dce21 100644 --- a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py +++ b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py @@ -18,20 +18,24 @@ from pyhugegraph.client import PyHugeClient if __name__ == "__main__": - client = PyHugeClient( - url="http://127.0.0.1:8080", user="admin", pwd="admin", graph="hugegraph", graphspace=None - ) + client = PyHugeClient(url="http://127.0.0.1:8080", user="admin", pwd="admin", graph="hugegraph", graphspace=None) """schema""" schema = client.schema() schema.propertyKey("name").asText().ifNotExist().create() schema.propertyKey("birthDate").asText().ifNotExist().create() schema.vertexLabel("Person").properties("name", "birthDate").usePrimaryKeyId().primaryKeys( +<<<<<<< HEAD "name" ).ifNotExist().create() schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys( "name" ).ifNotExist().create() +======= + "name" + ).ifNotExist().create() + schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys("name").ifNotExist().create() +>>>>>>> 87ee5d3 (style: format code with black line-length 120) schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() print(schema.getVertexLabels()) diff --git a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_test.py b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_test.py index 2bfe6ea97..075d069b7 100644 --- a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_test.py +++ b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_test.py @@ -31,17 +31,14 @@ def __init__( from pyhugegraph.client import PyHugeClient except ImportError: raise ValueError( - "Please install HugeGraph Python client first: " - "`pip3 install hugegraph-python-client`" + "Please install HugeGraph Python client first: " "`pip3 install hugegraph-python-client`" ) from ImportError self.username = username self.password = password self.url = url self.graph = graph - self.client = PyHugeClient( - url=url, user=username, pwd=password, graph=graph, graphspace=None - ) + self.client = PyHugeClient(url=url, user=username, pwd=password, graph=graph, graphspace=None) self.schema = "" def exec(self, query) -> str: diff --git a/hugegraph-python-client/src/pyhugegraph/structure/vertex_label_data.py b/hugegraph-python-client/src/pyhugegraph/structure/vertex_label_data.py index 39da4eebb..aaee1370f 100644 --- a/hugegraph-python-client/src/pyhugegraph/structure/vertex_label_data.py +++ b/hugegraph-python-client/src/pyhugegraph/structure/vertex_label_data.py @@ -65,8 +65,5 @@ def enableLabelIndex(self): return self.__enable_label_index def __repr__(self): - res = ( - f"name: {self.__name}, primary_keys: {self.__primary_keys}, " - f"properties: {self.__properties}" - ) + res = f"name: {self.__name}, primary_keys: {self.__primary_keys}, " f"properties: {self.__properties}" return res diff --git a/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py b/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py index f4a38a418..7acf2fdfa 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py @@ -145,9 +145,7 @@ def wrapper(self: "HGraphContext", *args: Any, **kwargs: Any) -> Any: class RouterMixin: - def _invoke_request_registered( - self, placeholders: dict = None, validator=ResponseValidation(), **kwargs: Any - ): + def _invoke_request_registered(self, placeholders: dict = None, validator=ResponseValidation(), **kwargs: Any): """ Make an HTTP request using the stored partial request function. Args: diff --git a/hugegraph-python-client/src/pyhugegraph/utils/util.py b/hugegraph-python-client/src/pyhugegraph/utils/util.py index 56a135547..6df64deff 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/util.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/util.py @@ -34,8 +34,7 @@ def create_exception(response_content): data = json.loads(response_content) if "ServiceUnavailableException" in data.get("exception", ""): raise ServiceUnavailableException( - f'ServiceUnavailableException, "message": "{data["message"]}",' - f' "cause": "{data["cause"]}"' + f'ServiceUnavailableException, "message": "{data["message"]}",' f' "cause": "{data["cause"]}"' ) except (json.JSONDecodeError, KeyError) as e: raise Exception(f"Error parsing response content: {response_content}") from e @@ -59,10 +58,14 @@ def check_if_success(response, error=None): req_body = req.body if req.body else "Empty body" response_body = response.text if response.text else "Empty body" log.error( +<<<<<<< HEAD "Error-Client: Request URL: %s, Request Body: %s, Response Body: %s", req.url, req_body, response_body, +======= + "Error-Client: Request URL: %s, Request Body: %s, Response Body: %s", req.url, req_body, response_body +>>>>>>> 87ee5d3 (style: format code with black line-length 120) ) raise error return True diff --git a/hugegraph-python-client/src/tests/api/test_traverser.py b/hugegraph-python-client/src/tests/api/test_traverser.py index 70c206acc..ae44cf6f8 100644 --- a/hugegraph-python-client/src/tests/api/test_traverser.py +++ b/hugegraph-python-client/src/tests/api/test_traverser.py @@ -55,9 +55,7 @@ def test_traverser_operations(self): self.assertEqual(k_out_result["vertices"], ["1:peter", "2:ripple"]) k_neighbor_result = self.traverser.k_neighbor(marko, 2) - self.assertEqual( - k_neighbor_result["vertices"], ["1:peter", "1:josh", "2:lop", "2:ripple", "1:vadas"] - ) + self.assertEqual(k_neighbor_result["vertices"], ["1:peter", "1:josh", "2:lop", "2:ripple", "1:vadas"]) same_neighbors_result = self.traverser.same_neighbors(marko, josh) self.assertEqual(same_neighbors_result["same_neighbors"], ["2:lop"]) @@ -69,16 +67,10 @@ def test_traverser_operations(self): self.assertEqual(shortest_path_result["path"], ["1:marko", "1:josh", "2:ripple"]) all_shortest_paths_result = self.traverser.all_shortest_paths(marko, ripple, 3) - self.assertEqual( - all_shortest_paths_result["paths"], [{"objects": ["1:marko", "1:josh", "2:ripple"]}] - ) + self.assertEqual(all_shortest_paths_result["paths"], [{"objects": ["1:marko", "1:josh", "2:ripple"]}]) - weighted_shortest_path_result = self.traverser.weighted_shortest_path( - marko, ripple, "weight", 3 - ) - self.assertEqual( - weighted_shortest_path_result["vertices"], ["1:marko", "1:josh", "2:ripple"] - ) + weighted_shortest_path_result = self.traverser.weighted_shortest_path(marko, ripple, "weight", 3) + self.assertEqual(weighted_shortest_path_result["vertices"], ["1:marko", "1:josh", "2:ripple"]) single_source_shortest_path_result = self.traverser.single_source_shortest_path(marko, 2) self.assertEqual( @@ -92,9 +84,7 @@ def test_traverser_operations(self): }, ) - multi_node_shortest_path_result = self.traverser.multi_node_shortest_path( - [marko, josh], max_depth=2 - ) + multi_node_shortest_path_result = self.traverser.multi_node_shortest_path([marko, josh], max_depth=2) self.assertEqual( multi_node_shortest_path_result["vertices"], [ @@ -131,9 +121,7 @@ def test_traverser_operations(self): } ], ) - self.assertEqual( - customized_paths_result["paths"], [{"objects": ["1:marko", "2:lop"], "weights": [8.0]}] - ) + self.assertEqual(customized_paths_result["paths"], [{"objects": ["1:marko", "2:lop"], "weights": [8.0]}]) sources = {"ids": [], "label": "person", "properties": {"name": "vadas"}} @@ -186,15 +174,11 @@ def test_traverser_operations(self): sources = {"ids": ["2:lop", "2:ripple"]} path_patterns = [{"steps": [{"direction": "IN", "labels": ["created"], "max_degree": -1}]}] - customized_crosspoints_result = self.traverser.customized_crosspoints( - sources, path_patterns - ) + customized_crosspoints_result = self.traverser.customized_crosspoints(sources, path_patterns) self.assertEqual(customized_crosspoints_result["crosspoints"], ["1:josh"]) rings_result = self.traverser.rings(marko, 3) - self.assertEqual( - rings_result["rings"], [{"objects": ["1:marko", "2:lop", "1:josh", "1:marko"]}] - ) + self.assertEqual(rings_result["rings"], [{"objects": ["1:marko", "2:lop", "1:josh", "1:marko"]}]) rays_result = self.traverser.rays(marko, 2) self.assertEqual( diff --git a/hugegraph-python-client/src/tests/client_utils.py b/hugegraph-python-client/src/tests/client_utils.py index f711072b8..11cbb4a55 100644 --- a/hugegraph-python-client/src/tests/client_utils.py +++ b/hugegraph-python-client/src/tests/client_utils.py @@ -59,23 +59,21 @@ def init_property_key(self): def init_vertex_label(self): schema = self.schema - schema.vertexLabel("person").properties("name", "age", "city").primaryKeys( - "name" - ).nullableKeys("city").ifNotExist().create() - schema.vertexLabel("software").properties("name", "lang", "price").primaryKeys( - "name" - ).nullableKeys("price").ifNotExist().create() + schema.vertexLabel("person").properties("name", "age", "city").primaryKeys("name").nullableKeys( + "city" + ).ifNotExist().create() + schema.vertexLabel("software").properties("name", "lang", "price").primaryKeys("name").nullableKeys( + "price" + ).ifNotExist().create() schema.vertexLabel("book").useCustomizeStringId().properties("name", "price").nullableKeys( "price" ).ifNotExist().create() def init_edge_label(self): schema = self.schema - schema.edgeLabel("knows").sourceLabel("person").targetLabel( - "person" - ).multiTimes().properties("date", "city").sortKeys("date").nullableKeys( - "city" - ).ifNotExist().create() + schema.edgeLabel("knows").sourceLabel("person").targetLabel("person").multiTimes().properties( + "date", "city" + ).sortKeys("date").nullableKeys("city").ifNotExist().create() schema.edgeLabel("created").sourceLabel("person").targetLabel("software").properties( "date", "city" ).nullableKeys("city").ifNotExist().create() @@ -84,16 +82,10 @@ def init_index_label(self): schema = self.schema schema.indexLabel("personByCity").onV("person").by("city").secondary().ifNotExist().create() schema.indexLabel("personByAge").onV("person").by("age").range().ifNotExist().create() - schema.indexLabel("softwareByPrice").onV("software").by( - "price" - ).range().ifNotExist().create() - schema.indexLabel("softwareByLang").onV("software").by( - "lang" - ).secondary().ifNotExist().create() + schema.indexLabel("softwareByPrice").onV("software").by("price").range().ifNotExist().create() + schema.indexLabel("softwareByLang").onV("software").by("lang").secondary().ifNotExist().create() schema.indexLabel("knowsByDate").onE("knows").by("date").secondary().ifNotExist().create() - schema.indexLabel("createdByDate").onE("created").by( - "date" - ).secondary().ifNotExist().create() + schema.indexLabel("createdByDate").onE("created").by("date").secondary().ifNotExist().create() def init_vertices(self): graph = self.graph From 8f3ba72099fd7f09dadceaa3701d38fc8d1df396 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Tue, 12 Aug 2025 10:45:20 +0800 Subject: [PATCH 17/71] fix(security): add URL validation to avoid potential SSRF in test_api_connection --- .../src/hugegraph_llm/demo/rag_demo/configs_block.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index a98c636b9..ef2786e20 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -21,6 +21,7 @@ from typing import Optional import gradio as gr +import urllib.parse as _urlparse import requests from dotenv import dotenv_values from requests.auth import HTTPBasicAuth @@ -33,6 +34,15 @@ current_llm = "chat" +def _validate_url_safe(url: str) -> None: + """Basic SSRF guard: allow only http/https and forbid unexpected schemes.""" + parsed = _urlparse.urlparse(url) + if parsed.scheme not in {"http", "https"}: + raise gr.Error("Only http/https URLs are allowed for connection test.") + if not parsed.netloc: + raise gr.Error("URL missing hostname.") + + def test_litellm_embedding(api_key, api_base, model_name, model_dim) -> int: llm_client = LiteLLMEmbedding( <<<<<<< HEAD @@ -71,6 +81,7 @@ def test_litellm_chat(api_key, api_base, model_name, max_tokens: int) -> int: def test_api_connection(url, method="GET", headers=None, params=None, body=None, auth=None, origin_call=None) -> int: # TODO: use fastapi.request / starlette instead? log.debug("Request URL: %s", url) + _validate_url_safe(url) try: if method.upper() == "GET": resp = requests.get(url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth) From f90c1d56d5434c9c0cd36437d60921563b177e35 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Tue, 12 Aug 2025 14:45:20 +0800 Subject: [PATCH 18/71] small fix --- hugegraph-llm/pyproject.toml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index e5e1f502d..ca7ee3914 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -99,10 +99,6 @@ line-length = 120 indent-width = 4 extend-exclude = [] extend-select = ["I"] -<<<<<<< HEAD -======= -extend-select = ["I"] ->>>>>>> 56bcac3 (feat(llm): import sort && change name) [tool.ruff.format] quote-style = "preserve" From 3561876d4135df01ff31edd9904e3307f6da0ecb Mon Sep 17 00:00:00 2001 From: lingxiao Date: Tue, 12 Aug 2025 14:53:31 +0800 Subject: [PATCH 19/71] fix url --- .../demo/rag_demo/configs_block.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index ef2786e20..54aac54a0 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -42,6 +42,32 @@ def _validate_url_safe(url: str) -> None: if not parsed.netloc: raise gr.Error("URL missing hostname.") + # 防止 SSRF：禁止访问本地主机或内网地址 + hostname = parsed.hostname or "" + # IPv4 私有网段 + private_ipv4_networks = [ + ("10.",), + ("172.", range(16, 32)), + ("192.168.",), + ("127.",), + ("0.",), + ] + + def _is_private_ipv4(host: str) -> bool: + for prefix in private_ipv4_networks: + base = prefix[0] + if host.startswith(base): + # 处理 172.16.0.0/12 特例 + if len(prefix) == 1: + return True + if int(host.split(".")[1]) in prefix[1]: + return True + return False + + # IPv6 localhost + if hostname in {"localhost", "::1"} or _is_private_ipv4(hostname): + raise gr.Error("Connection to localhost or private network addresses is not allowed.") + def test_litellm_embedding(api_key, api_base, model_name, model_dim) -> int: llm_client = LiteLLMEmbedding( From a1c128e0c1a9d1ebff4bd985ad75217dfe05f8c0 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Wed, 13 Aug 2025 18:36:27 +0800 Subject: [PATCH 20/71] fix --- hugegraph-llm/src/hugegraph_llm/api/admin_api.py | 1 - hugegraph-llm/src/hugegraph_llm/config/llm_config.py | 9 ++++----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py index 9bf2ff275..96db7da0a 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py @@ -33,7 +33,6 @@ async def log_stream_api(req: LogStreamRequest): if admin_settings.admin_token != req.admin_token: raise generate_response( RAGResponse(status_code=status.HTTP_403_FORBIDDEN, message="Invalid admin_token") # pylint: disable=E0702 - ) ) log_path = os.path.join("logs", req.log_file) diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index 493eca287..a63d401c1 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -63,11 +63,10 @@ class LLMConfig(BaseConfig): ollama_text2gql_port: int = 11434 ollama_text2gql_language_model: str | None = None ollama_embedding_host: str = "127.0.0.1" - ollama_embedding_port: int = int(11434) - ollama_embedding_model: str = 'quentinz/bge-large-zh-v1.5' - ollama_embedding_model_dim: Optional[int] = ( - int(os.getenv("OLLAMA_EMBEDDING_MODEL_DIM")) if os.getenv("OLLAMA_EMBEDDING_MODEL_DIM") else None # type:ignore - ) + ollama_embedding_port: int = 11434 + ollama_embedding_model: str = "quentinz/bge-large-zh-v1.5" + _env_ollama_dim = os.getenv("OLLAMA_EMBEDDING_MODEL_DIM") + ollama_embedding_model_dim: Optional[int] = int(_env_ollama_dim) if _env_ollama_dim else None # 4. QianFan/WenXin settings # TODO: update to one token key mode From 7eedb1837816a94d5d89412d88e6941e2ba61a36 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Wed, 13 Aug 2025 18:40:17 +0800 Subject: [PATCH 21/71] fix --- hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py | 2 +- .../src/hugegraph_llm/operators/common_op/check_schema.py | 2 +- hugegraph-ml/src/hugegraph_ml/models/bgrl.py | 2 +- hugegraph-python-client/src/pyhugegraph/api/gremlin.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 54aac54a0..d27ea7b4d 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -20,8 +20,8 @@ from functools import partial from typing import Optional -import gradio as gr import urllib.parse as _urlparse +import gradio as gr import requests from dotenv import dotenv_values from requests.auth import HTTPBasicAuth diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index fc729c11e..47b0f060f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -16,7 +16,7 @@ # under the License. -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional from hugegraph_llm.enums.property_cardinality import PropertyCardinality from hugegraph_llm.enums.property_data_type import PropertyDataType diff --git a/hugegraph-ml/src/hugegraph_ml/models/bgrl.py b/hugegraph-ml/src/hugegraph_ml/models/bgrl.py index 288a434b8..0000e546c 100644 --- a/hugegraph-ml/src/hugegraph_ml/models/bgrl.py +++ b/hugegraph-ml/src/hugegraph_ml/models/bgrl.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -# pylint: disable=C0103,R1705.R1734 +# pylint: disable=C0103,R1705.R1734,E1102 """ Bootstrapped Graph Latents (BGRL) diff --git a/hugegraph-python-client/src/pyhugegraph/api/gremlin.py b/hugegraph-python-client/src/pyhugegraph/api/gremlin.py index d91bdb04b..3261d60b3 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/gremlin.py +++ b/hugegraph-python-client/src/pyhugegraph/api/gremlin.py @@ -43,7 +43,7 @@ def exec(self, gremlin): try: if response := self._invoke_request(data=gremlin_data.to_json()): return ResponseData(response).result - log.error(f"Gremlin can't get results: {str(response)}") # pylint: disable=logging-fstring-interpolation + log.error("Gremlin can't get results: %s", str(response)) return None except Exception as e: raise NotFoundError(f"Gremlin can't get results: {e}") from e From 68e06fc23936bc5c9ed51f63a1ee029a8019edd7 Mon Sep 17 00:00:00 2001 From: imbajin Date: Wed, 27 Aug 2025 21:05:08 +0800 Subject: [PATCH 22/71] chore: mark vectordb optional --- hugegraph-llm/README.md | 3 + hugegraph-llm/pyproject.toml | 9 +++ .../hugegraph_llm/api/models/rag_requests.py | 64 +++------------ .../src/hugegraph_llm/api/rag_api.py | 8 +- .../hugegraph_llm/config/models/__init__.py | 2 + .../demo/rag_demo/admin_block.py | 4 + .../demo/rag_demo/configs_block.py | 22 ++++- .../hugegraph_llm/demo/rag_demo/rag_block.py | 25 +++--- .../demo/rag_demo/text2gremlin_block.py | 19 ++--- .../demo/rag_demo/vector_graph_block.py | 4 + .../src/hugegraph_llm/document/chunk_split.py | 6 ++ .../vector_index/faiss_vector_store.py | 4 +- .../src/hugegraph_llm/models/llms/ollama.py | 4 + .../hugegraph_llm/models/rerankers/cohere.py | 8 +- .../models/rerankers/init_reranker.py | 4 +- .../models/rerankers/siliconflow.py | 8 +- .../operators/common_op/check_schema.py | 22 ++--- .../operators/common_op/merge_dedup_rerank.py | 11 +-- .../hugegraph_llm/operators/graph_rag_task.py | 4 + .../hugegraph_op/commit_to_hugegraph.py | 20 ++--- .../hugegraph_op/fetch_graph_data.py | 1 - .../operators/hugegraph_op/graph_rag_query.py | 50 ++++-------- .../operators/hugegraph_op/schema_manager.py | 8 +- .../index_op/build_semantic_index.py | 5 ++ .../index_op/gremlin_example_index_query.py | 3 + .../operators/index_op/vector_index_query.py | 4 + .../operators/llm_op/answer_synthesize.py | 80 ++++++++++--------- .../operators/llm_op/disambiguate_data.py | 4 + .../operators/llm_op/gremlin_generate.py | 8 +- .../operators/llm_op/info_extract.py | 12 +++ .../operators/llm_op/keyword_extract.py | 8 +- .../llm_op/property_graph_extract.py | 14 ++-- .../src/hugegraph_llm/utils/anchor.py | 3 +- .../hugegraph_llm/utils/embedding_utils.py | 4 +- .../hugegraph_llm/utils/hugegraph_utils.py | 12 +-- 35 files changed, 219 insertions(+), 248 deletions(-) diff --git a/hugegraph-llm/README.md b/hugegraph-llm/README.md index 526320d4a..e0ebfdc63 100644 --- a/hugegraph-llm/README.md +++ b/hugegraph-llm/README.md @@ -113,6 +113,9 @@ python -m hugegraph_llm.demo.rag_demo.app --host 127.0.0.1 --port 18001 > The following commands assume you're in the activated virtual environment from step 4 above ```bash +# To use vector database backends (e.g., Milvus, Qdrant), sync the optional dependencies: +uv sync --extra vectordb + # Download NLTK stopwords for better text processing python ./src/hugegraph_llm/operators/common_op/nltk_helper.py diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index ca7ee3914..0d12210fb 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -61,6 +61,12 @@ dependencies = [ "hugegraph-python-client", "pycgraph", ] + +[project.optional-dependencies] +vectordb = [ + "pymilvus==2.5.9", + "qdrant-client==1.14.2", +] [project.urls] homepage = "https://hugegraph.apache.org/" repository = "https://github.com/apache/incubator-hugegraph-ai" @@ -98,6 +104,9 @@ disallow_untyped_defs = false line-length = 120 indent-width = 4 extend-exclude = [] + +# TODO: move this config in the root pyproject.toml & add more rules for it +[tool.ruff.lint] extend-select = ["I"] [tool.ruff.format] diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py index 720c25a86..5222f0cfa 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py +++ b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py @@ -36,23 +36,13 @@ class RAGRequest(BaseModel): raw_answer: bool = Query(False, description="Use LLM to generate answer directly") vector_only: bool = Query(False, description="Use LLM to generate answer with vector") graph_only: bool = Query(True, description="Use LLM to generate answer with graph RAG only") - graph_vector_answer: bool = Query( - False, description="Use LLM to generate answer with vector & GraphRAG" - ) + graph_vector_answer: bool = Query(False, description="Use LLM to generate answer with vector & GraphRAG") graph_ratio: float = Query(0.5, description="The ratio of GraphRAG ans & vector ans") - rerank_method: Literal["bleu", "reranker"] = Query( - "bleu", description="Method to rerank the results." - ) - near_neighbor_first: bool = Query( - False, description="Prioritize near neighbors in the search results." - ) - custom_priority_info: str = Query( - "", description="Custom information to prioritize certain results." - ) + rerank_method: Literal["bleu", "reranker"] = Query("bleu", description="Method to rerank the results.") + near_neighbor_first: bool = Query(False, description="Prioritize near neighbors in the search results.") + custom_priority_info: str = Query("", description="Custom information to prioritize certain results.") # Graph Configs - max_graph_items: int = Query( - 30, description="Maximum number of items for GQL queries in graph." - ) + max_graph_items: int = Query(30, description="Maximum number of items for GQL queries in graph.") topk_return_results: int = Query(20, description="Number of sorted results to return finally.") vector_dis_threshold: float = Query( 0.9, @@ -64,18 +54,10 @@ class RAGRequest(BaseModel): description="TopK results returned for each keyword \ extracted from the query, by default only the most similar one is returned.", ) -<<<<<<< HEAD - client_config: Optional[GraphConfigRequest] = Query( - None, description="hugegraph server config." - ) -======= client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") ->>>>>>> 87ee5d3 (style: format code with black line-length 120) # Keep prompt params in the end - answer_prompt: Optional[str] = Query( - prompt.answer_prompt, description="Prompt to guide the answer generation." - ) + answer_prompt: Optional[str] = Query(prompt.answer_prompt, description="Prompt to guide the answer generation.") keywords_extract_prompt: Optional[str] = Query( prompt.keywords_extract_prompt, description="Prompt for extracting keywords from query.", @@ -91,9 +73,7 @@ class RAGRequest(BaseModel): class GraphRAGRequest(BaseModel): query: str = Query(..., description="Query you want to ask") # Graph Configs - max_graph_items: int = Query( - 30, description="Maximum number of items for GQL queries in graph." - ) + max_graph_items: int = Query(30, description="Maximum number of items for GQL queries in graph.") topk_return_results: int = Query(20, description="Number of sorted results to return finally.") vector_dis_threshold: float = Query( 0.9, @@ -106,24 +86,16 @@ class GraphRAGRequest(BaseModel): from the query, by default only the most similar one is returned.", ) - client_config: Optional[GraphConfigRequest] = Query( - None, description="hugegraph server config." - ) + client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") get_vertex_only: bool = Query(False, description="return only keywords & vertex (early stop).") gremlin_tmpl_num: int = Query( 1, description="Number of Gremlin templates to use. If num <=0 means template is not provided", ) - rerank_method: Literal["bleu", "reranker"] = Query( - "bleu", description="Method to rerank the results." - ) - near_neighbor_first: bool = Query( - False, description="Prioritize near neighbors in the search results." - ) - custom_priority_info: str = Query( - "", description="Custom information to prioritize certain results." - ) + rerank_method: Literal["bleu", "reranker"] = Query("bleu", description="Method to rerank the results.") + near_neighbor_first: bool = Query(False, description="Prioritize near neighbors in the search results.") + custom_priority_info: str = Query("", description="Custom information to prioritize certain results.") gremlin_prompt: Optional[str] = Query( prompt.gremlin_generate_prompt, description="Prompt for the Text2Gremlin query.", @@ -167,20 +139,12 @@ class GremlinOutputType(str, Enum): class GremlinGenerateRequest(BaseModel): query: str -<<<<<<< HEAD - example_num: Optional[int] = Query( - 0, description="Number of Gremlin templates to use.(0 means no templates)" - ) -======= example_num: Optional[int] = Query(0, description="Number of Gremlin templates to use.(0 means no templates)") ->>>>>>> 87ee5d3 (style: format code with black line-length 120) gremlin_prompt: Optional[str] = Query( prompt.gremlin_generate_prompt, description="Prompt for the Text2Gremlin query.", ) - client_config: Optional[GraphConfigRequest] = Query( - None, description="hugegraph server config." - ) + client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") output_types: Optional[List[GremlinOutputType]] = Query( default=[GremlinOutputType.TEMPLATE_GREMLIN], description=""" @@ -197,7 +161,5 @@ def validate_prompt_placeholders(cls, v): required_placeholders = ["{query}", "{schema}", "{example}", "{vertices}"] missing = [p for p in required_placeholders if p not in v] if missing: - raise ValueError( - f"Prompt template is missing required placeholders: {', '.join(missing)}" - ) + raise ValueError(f"Prompt template is missing required placeholders: {', '.join(missing)}") return v diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py index 5c9295efa..bcd15f54f 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py @@ -167,9 +167,7 @@ def embedding_config_api(req: LLMConfigRequest): llm_settings.embedding_type = req.llm_type if req.llm_type == "openai": - res = apply_embedding_conf( - req.api_key, req.api_base, req.language_model, origin_call="http" - ) + res = apply_embedding_conf(req.api_key, req.api_base, req.language_model, origin_call="http") else: res = apply_embedding_conf(req.host, req.port, req.language_model, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) @@ -179,9 +177,7 @@ def rerank_config_api(req: RerankerConfigRequest): llm_settings.reranker_type = req.reranker_type if req.reranker_type == "cohere": - res = apply_reranker_conf( - req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http" - ) + res = apply_reranker_conf(req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http") elif req.reranker_type == "siliconflow": res = apply_reranker_conf(req.api_key, req.reranker_model, None, origin_call="http") else: diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py b/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py index e73646fd1..d7738036e 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py @@ -17,3 +17,5 @@ from .base_config import BaseConfig from .base_prompt_config import BasePromptConfig + +__all__ = ["BaseConfig", "BasePromptConfig"] \ No newline at end of file diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index ec01741e3..78ac89ad9 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -109,6 +109,7 @@ def create_admin_block(): ) # Error message box, initially hidden +<<<<<<< HEAD <<<<<<< HEAD error_message = gr.Textbox( label="", visible=False, interactive=False, elem_classes="error-message" @@ -116,6 +117,9 @@ def create_admin_block(): ======= error_message = gr.Textbox(label="", visible=False, interactive=False, elem_classes="error-message") >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + error_message = gr.Textbox(label="", visible=False, interactive=False, elem_classes="error-message") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) # Button to submit password submit_button = gr.Button("Submit") diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index d27ea7b4d..d43bee8b8 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -17,10 +17,10 @@ import json import os +import urllib.parse as _urlparse from functools import partial from typing import Optional -import urllib.parse as _urlparse import gradio as gr import requests from dotenv import dotenv_values @@ -177,6 +177,7 @@ def apply_embedding_config(arg1, arg2, arg3, arg4, origin_call=None) -> int: test_url = llm_settings.openai_embedding_api_base + "/embeddings" headers = {"Authorization": f"Bearer {arg1}"} data = {"model": arg3, "input": "test"} +<<<<<<< HEAD <<<<<<< HEAD status_code = test_api_connection( test_url, method="POST", headers=headers, body=data, origin_call=origin_call @@ -188,6 +189,9 @@ def apply_embedding_config(arg1, arg2, arg3, arg4, origin_call=None) -> int: llm_settings.qianfan_embedding_model = arg3 llm_settings.qianfan_embedding_model_dim = arg4 >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) +>>>>>>> 8e0bf08 (chore: mark vectordb optional) elif embedding_option == "ollama/local": llm_settings.ollama_embedding_host = arg1 llm_settings.ollama_embedding_port = int(arg2) @@ -284,6 +288,7 @@ def apply_llm_config( setattr(llm_settings, f"openai_{current_llm_config}_language_model", model_name) setattr(llm_settings, f"openai_{current_llm_config}_tokens", int(max_tokens)) +<<<<<<< HEAD <<<<<<< HEAD test_url = ( getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" @@ -291,6 +296,9 @@ def apply_llm_config( ======= test_url = getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + test_url = getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" +>>>>>>> 8e0bf08 (chore: mark vectordb optional) data = { "model": model_name, "temperature": 0.01, @@ -298,6 +306,7 @@ def apply_llm_config( } <<<<<<< HEAD headers = {"Authorization": f"Bearer {api_key_or_host}"} +<<<<<<< HEAD <<<<<<< HEAD status_code = test_api_connection( test_url, method="POST", headers=headers, body=data, origin_call=origin_call @@ -314,6 +323,9 @@ def apply_llm_config( ======= status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) +>>>>>>> 8e0bf08 (chore: mark vectordb optional) elif llm_option == "ollama/local": setattr(llm_settings, f"ollama_{current_llm_config}_host", api_key_or_host) @@ -390,9 +402,7 @@ def create_configs_block() -> list: ), ] graph_config_button = gr.Button("Apply Configuration") - graph_config_button.click( - apply_graph_config, inputs=graph_config_input - ) # pylint: disable=no-member + graph_config_button.click(apply_graph_config, inputs=graph_config_input) # pylint: disable=no-member # TODO : use OOP to refactor the following code with gr.Accordion("2. Set up the LLM.", open=False): @@ -548,10 +558,14 @@ def chat_llm_settings(llm_type): if not api_extract_key: llm_config_button.click(apply_llm_config_with_text2gql_op, inputs=llm_config_input) if not api_text2sql_key: +<<<<<<< HEAD <<<<<<< HEAD llm_config_button.click( apply_llm_config_with_extract_op, inputs=llm_config_input ) +======= + llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) +>>>>>>> 8e0bf08 (chore: mark vectordb optional) with gr.Tab(label="mini_tasks"): extract_llm_dropdown = gr.Dropdown( diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index cf40c444e..396727f82 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -349,19 +349,20 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): +<<<<<<< HEAD raw_radio = gr.Radio( choices=[True, False], value=False, label="Basic LLM Answer" ) vector_only_radio = gr.Radio( choices=[True, False], value=False, label="Vector-only Answer" ) +======= + raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") + vector_only_radio = gr.Radio(choices=[True, False], value=False, label="Vector-only Answer") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) with gr.Row(): - graph_only_radio = gr.Radio( - choices=[True, False], value=True, label="Graph-only Answer" - ) - graph_vector_radio = gr.Radio( - choices=[True, False], value=False, label="Graph-Vector Answer" - ) + graph_only_radio = gr.Radio(choices=[True, False], value=True, label="Graph-only Answer") + graph_vector_radio = gr.Radio(choices=[True, False], value=False, label="Graph-Vector Answer") def toggle_slider(enable): return gr.update(interactive=enable) @@ -379,13 +380,9 @@ def toggle_slider(enable): label="Template Num (<0 means disable text2gql) ", precision=0, ) - graph_ratio = gr.Slider( - 0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False - ) + graph_ratio = gr.Slider(0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False) - graph_vector_radio.change( - toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio - ) # pylint: disable=no-member + graph_vector_radio.change(toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio) # pylint: disable=no-member near_neighbor_first = gr.Checkbox( value=False, label="Near neighbor first(Optional)", @@ -514,9 +511,7 @@ def several_rag_answer( with gr.Row(): with gr.Column(): - questions_file = gr.File( - file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" - ) + questions_file = gr.File(file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)") with gr.Column(): test_template_file = os.path.join( resource_path, "demo", "questions_template.xlsx" diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 32529fb1e..fe885bfa9 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -219,9 +219,14 @@ def _execute_queries(context, output_types): def gremlin_generate( inp, example_num, schema, gremlin_prompt, requested_outputs: Optional[List[str]] = None ) -> GremlinResult: +<<<<<<< HEAD generator = GremlinGenerator( llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() ) +======= + vector_index = get_vector_index_class(index_settings.now_vector_index) + generator = GremlinGenerator(llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding()) +>>>>>>> 8e0bf08 (chore: mark vectordb optional) sm = SchemaManager(graph_name=schema) processed_schema, short_schema = _process_schema(schema, generator, sm) @@ -241,9 +246,7 @@ def gremlin_generate( _execute_queries(context, output_types) - match_result = json.dumps( - context.get("match_result", "No Results"), ensure_ascii=False, indent=2 - ) + match_result = json.dumps(context.get("match_result", "No Results"), ensure_ascii=False, indent=2) return GremlinResult.success_result( match_result=match_result, template_gremlin=context["result"], @@ -267,11 +270,7 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: if "edgelabels" in schema: mini_schema["edgelabels"] = [] for edge in schema["edgelabels"]: - new_edge = { - key: edge[key] - for key in ["name", "source_label", "target_label", "properties"] - if key in edge - } + new_edge = {key: edge[key] for key in ["name", "source_label", "target_label", "properties"] if key in edge} mini_schema["edgelabels"].append(new_edge) return mini_schema @@ -369,9 +368,7 @@ def create_text2gremlin_block() -> Tuple: ) with gr.Column(scale=1): - example_num_slider = gr.Slider( - minimum=0, maximum=10, step=1, value=2, label="Number of refer examples" - ) + example_num_slider = gr.Slider(minimum=0, maximum=10, step=1, value=2, label="Number of refer examples") schema_box = gr.Textbox( value=prompt.text2gql_graph_schema, label="Schema", lines=2, show_copy_button=True ) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 0efc11d59..af6570b19 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -100,6 +100,7 @@ def load_query_examples(): language = getattr( prompt, "language", +<<<<<<< HEAD <<<<<<< HEAD ( getattr(prompt.llm_settings, "language", "EN") @@ -109,6 +110,9 @@ def load_query_examples(): ======= getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") else "EN", >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + (getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") else "EN"), +>>>>>>> 8e0bf08 (chore: mark vectordb optional) ) if language.upper() == "CN": examples_path = os.path.join(resource_path, "prompt_examples", "query_examples_CN.json") diff --git a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py index 12572a9e7..5012f7e01 100644 --- a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py +++ b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py @@ -37,6 +37,7 @@ def __init__( else: raise ValueError("Argument `language` must be zh or en!") if split_type == "paragraph": +<<<<<<< HEAD <<<<<<< HEAD self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=500, chunk_overlap=30, separators=separators @@ -50,6 +51,11 @@ def __init__( elif split_type == "sentence": self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=50, chunk_overlap=0, separators=separators) >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=30, separators=separators) + elif split_type == "sentence": + self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=50, chunk_overlap=0, separators=separators) +>>>>>>> 8e0bf08 (chore: mark vectordb optional) else: raise ValueError("Arg `type` must be paragraph, sentence!") diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py index a8f23155a..d0a016c53 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py @@ -66,9 +66,7 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: self.properties = [p for i, p in enumerate(self.properties) if i not in indices] return remove_num - def search( - self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9 - ) -> List[Any]: + def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9) -> List[Any]: if self.index.ntotal == 0: return [] diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py index 1ca304baa..2fa979ed9 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py @@ -118,6 +118,7 @@ async def agenerate_streaming( messages = [{"role": "user", "content": prompt}] try: +<<<<<<< HEAD <<<<<<< HEAD async_generator = await self.async_client.chat( model=self.model, messages=messages, stream=True @@ -125,6 +126,9 @@ async def agenerate_streaming( ======= async_generator = await self.async_client.chat(model=self.model, messages=messages, stream=True) >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + async_generator = await self.async_client.chat(model=self.model, messages=messages, stream=True) +>>>>>>> 8e0bf08 (chore: mark vectordb optional) async for chunk in async_generator: token = chunk.get("message", {}).get("content", "") if on_token_callback: diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py index 3bf481ce2..fd4643e3c 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py @@ -31,14 +31,10 @@ def __init__( self.base_url = base_url self.model = model - def get_rerank_lists( - self, query: str, documents: List[str], top_n: Optional[int] = None - ) -> List[str]: + def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int] = None) -> List[str]: if not top_n: top_n = len(documents) - assert top_n <= len( - documents - ), "'top_n' should be less than or equal to the number of documents" + assert top_n <= len(documents), "'top_n' should be less than or equal to the number of documents" if top_n == 0: return [] diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py index 6136d61b4..aa9f0c061 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py @@ -32,7 +32,5 @@ def get_reranker(self): model=llm_settings.reranker_model, ) if self.reranker_type == "siliconflow": - return SiliconReranker( - api_key=llm_settings.reranker_api_key, model=llm_settings.reranker_model - ) + return SiliconReranker(api_key=llm_settings.reranker_api_key, model=llm_settings.reranker_model) raise Exception("Reranker type is not supported!") diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py index e4a9b550a..a67a6ef25 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py @@ -29,14 +29,10 @@ def __init__( self.api_key = api_key self.model = model - def get_rerank_lists( - self, query: str, documents: List[str], top_n: Optional[int] = None - ) -> List[str]: + def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int] = None) -> List[str]: if not top_n: top_n = len(documents) - assert top_n <= len( - documents - ), "'top_n' should be less than or equal to the number of documents" + assert top_n <= len(documents), "'top_n' should be less than or equal to the number of documents" if top_n == 0: return [] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index 47b0f060f..63618aaec 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -72,9 +72,7 @@ def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): property_label_set = {label["name"] for label in property_labels} return property_labels, property_label_set - def _process_vertex_labels( - self, schema: Dict[str, Any], property_labels: list, property_label_set: set - ) -> None: + def _process_vertex_labels(self, schema: Dict[str, Any], property_labels: list, property_label_set: set) -> None: for vertex_label in schema["vertexlabels"]: self._validate_vertex_label(vertex_label) properties = vertex_label["properties"] @@ -86,9 +84,7 @@ def _process_vertex_labels( vertex_label["nullable_keys"] = nullable_keys self._add_missing_properties(properties, property_labels, property_label_set) - def _process_edge_labels( - self, schema: Dict[str, Any], property_labels: list, property_label_set: set - ) -> None: + def _process_edge_labels(self, schema: Dict[str, Any], property_labels: list, property_label_set: set) -> None: for edge_label in schema["edgelabels"]: self._validate_edge_label(edge_label) properties = edge_label.get("properties", []) @@ -111,14 +107,8 @@ def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: check_type(edge_label, dict, "EdgeLabel in input data is not a dictionary.") - if ( - "name" not in edge_label - or "source_label" not in edge_label - or "target_label" not in edge_label - ): - log_and_raise( - "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." - ) + if "name" not in edge_label or "source_label" not in edge_label or "target_label" not in edge_label: + log_and_raise("EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'.") check_type(edge_label["name"], str, "'name' in edge_label is not of correct type.") check_type( edge_label["source_label"], @@ -137,9 +127,7 @@ def _process_keys(self, label: Dict[str, Any], key_type: str, default_keys: list new_keys = [key for key in keys if key in label["properties"]] return new_keys - def _add_missing_properties( - self, properties: list, property_labels: list, property_label_set: set - ) -> None: + def _add_missing_properties(self, properties: list, property_labels: list, property_label_set: set) -> None: for prop in properties: if prop not in property_label_set: property_labels.append( diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py index dc5b15e00..743cf3352 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py @@ -126,8 +126,7 @@ def _rerank_with_vertex_degree( reranker = Rerankers().get_reranker() try: vertex_rerank_res = [ - reranker.get_rerank_lists(query, vertex_degree) + [""] - for vertex_degree in vertex_degree_list + reranker.get_rerank_lists(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list ] except requests.exceptions.RequestException as e: log.warning( @@ -137,9 +136,7 @@ def _rerank_with_vertex_degree( self.switch_to_bleu = True if self.method == "bleu": - vertex_rerank_res = [ - _bleu_rerank(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list - ] + vertex_rerank_res = [_bleu_rerank(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list] depth = len(vertex_degree_list) for result in results: @@ -149,9 +146,7 @@ def _rerank_with_vertex_degree( knowledge_with_degree[result] += [""] * (depth - len(knowledge_with_degree[result])) def sort_key(res: str) -> Tuple[int, ...]: - return tuple( - vertex_rerank_res[i].index(knowledge_with_degree[res][i]) for i in range(depth) - ) + return tuple(vertex_rerank_res[i].index(knowledge_with_degree[res][i]) for i in range(depth)) sorted_results = sorted(results, key=sort_key) return sorted_results[:topn] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index f972c79b8..95eab701f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -251,6 +251,7 @@ def run(self, **kwargs) -> Dict[str, Any]: """ if len(self._operators) == 0: <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD self.extract_keywords().query_graphdb( max_graph_items=kwargs.get("max_graph_items") @@ -261,6 +262,9 @@ def run(self, **kwargs) -> Dict[str, Any]: ======= self.extract_keywords().query_graphdb(max_graph_items=kwargs.get("max_graph_items")).synthesize_answer() >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + self.extract_keywords().query_graphdb(max_graph_items=kwargs.get("max_graph_items")).synthesize_answer() +>>>>>>> 8e0bf08 (chore: mark vectordb optional) context = kwargs 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 52626b72b..7aefde57d 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 @@ -42,9 +42,7 @@ def run(self, data: dict) -> Dict[str, Any]: edges = data.get("edges", []) print(f"get schema {schema}") if not vertices and not edges: - log.critical( - "(Loading) Both vertices and edges are empty. Please check the input data again." - ) + log.critical("(Loading) Both vertices and edges are empty. Please check the input data again.") raise ValueError("Both vertices and edges input are empty.") if not schema: @@ -148,9 +146,7 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- continue # TODO: we could try batch add vertices first, setback to single-mode if failed - vid = self._handle_graph_creation( - self.client.graph().addVertex, input_label, input_properties - ).id + vid = self._handle_graph_creation(self.client.graph().addVertex, input_label, input_properties).id vertex["id"] = vid for edge in edges: @@ -197,19 +193,13 @@ def init_schema_if_need(self, schema: dict): def schema_free_mode(self, data): self.schema.propertyKey("name").asText().ifNotExist().create() - self.schema.vertexLabel("vertex").useCustomizeStringId().properties( - "name" - ).ifNotExist().create() + self.schema.vertexLabel("vertex").useCustomizeStringId().properties("name").ifNotExist().create() self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel("vertex").properties( "name" ).ifNotExist().create() - self.schema.indexLabel("vertexByName").onV("vertex").by( - "name" - ).secondary().ifNotExist().create() - self.schema.indexLabel("edgeByName").onE("edge").by( - "name" - ).secondary().ifNotExist().create() + 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: s, p, o = (element.strip() for element in item) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py index e93d916b3..4c4c167c4 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py @@ -22,7 +22,6 @@ class FetchGraphData: - def __init__(self, graph: PyHugeClient): self.graph = graph diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py index bcff5f07b..877a989bf 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py @@ -150,7 +150,7 @@ def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: if context["graph_result"]: context["graph_result_flag"] = 1 context["graph_context_head"] = ( - f"The following are graph query result " f"from gremlin query `{gremlin}`.\n" + f"The following are graph query result from gremlin query `{gremlin}`.\n" ) except Exception as e: # pylint: disable=broad-except log.error(e) @@ -196,8 +196,8 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: log.debug("Kneighbor gremlin query: %s", gremlin_query) paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( - self._format_graph_query_result(query_paths=paths) + graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = self._format_graph_query_result( + query_paths=paths ) # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) @@ -220,21 +220,17 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: max_deep=self._max_deep, max_items=self._max_items, ) - log.warning( - "Unable to find vid, downgraded to property query, please confirm if it meets expectation." - ) + log.warning("Unable to find vid, downgraded to property query, please confirm if it meets expectation.") paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( - self._format_graph_query_result(query_paths=paths) + graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = self._format_graph_query_result( + query_paths=paths ) context["graph_result"] = list(graph_chain_knowledge) if context["graph_result"]: context["graph_result_flag"] = 0 - context["vertex_degree_list"] = [ - list(vertex_degree) for vertex_degree in vertex_degree_list - ] + context["vertex_degree_list"] = [list(vertex_degree) for vertex_degree in vertex_degree_list] context["knowledge_with_degree"] = knowledge_with_degree context["graph_context_head"] = ( f"The following are graph knowledge in {self._max_deep} depth, e.g:\n" @@ -276,9 +272,7 @@ def _format_graph_from_vertex(self, query_result: List[Any]) -> Set[str]: knowledge.add(node_str) return knowledge - def _format_graph_query_result( - self, query_paths - ) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: + def _format_graph_query_result(self, query_paths) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: use_id_to_match = self._prop_to_match is None subgraph = set() subgraph_with_degree = {} @@ -288,9 +282,7 @@ def _format_graph_query_result( for path in query_paths: # 1. Process each path - path_str, vertex_with_degree = self._process_path( - path, use_id_to_match, v_cache, e_cache - ) + path_str, vertex_with_degree = self._process_path(path, use_id_to_match, v_cache, e_cache) subgraph.add(path_str) subgraph_with_degree[path_str] = vertex_with_degree # 2. Update vertex degree list @@ -352,9 +344,7 @@ def _process_vertex( return flat_rel, prior_edge_str_len, depth node_cache.add(matched_str) - props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v - ) + props_str = ", ".join(f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v) # TODO: we may remove label id or replace with label name if matched_str in v_cache: @@ -377,14 +367,10 @@ def _process_edge( use_id_to_match: bool, e_cache: Set[Tuple[str, str, str]], ) -> Tuple[str, int]: - props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v - ) + props_str = ", ".join(f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v) props_str = f"{{{props_str}}}" if props_str else "" prev_matched_str = ( - raw_flat_rel[i - 1]["id"] - if use_id_to_match - else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] + raw_flat_rel[i - 1]["id"] if use_id_to_match else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] ) edge_key = (item["inV"], item["label"], item["outV"]) @@ -394,16 +380,12 @@ def _process_edge( else: edge_label = item["label"] - edge_str = ( - f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" - ) + edge_str = f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" path_str += edge_str prior_edge_str_len = len(edge_str) return path_str, prior_edge_str_len - def _update_vertex_degree_list( - self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str] - ) -> None: + def _update_vertex_degree_list(self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str]) -> None: for depth, node_str in enumerate(nodes_with_degree): if depth >= len(vertex_degree_list): vertex_degree_list.append(set()) @@ -439,9 +421,7 @@ def _get_graph_schema(self, refresh: bool = False) -> str: relationships = schema.getRelations() self._schema = ( - f"Vertex properties: {vertex_schema}\n" - f"Edge properties: {edge_schema}\n" - f"Relationships: {relationships}\n" + f"Vertex properties: {vertex_schema}\nEdge properties: {edge_schema}\nRelationships: {relationships}\n" ) log.debug("Link(Relation): %s", relationships) return self._schema 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 2f0643a77..d8f59f50e 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 @@ -39,9 +39,7 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = { - key: vertex[key] for key in ["id", "name", "properties"] if key in vertex - } + new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) @@ -49,9 +47,7 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: mini_schema["edgelabels"] = [] for edge in schema["edgelabels"]: new_edge = { - key: edge[key] - for key in ["name", "source_label", "target_label", "properties"] - if key in edge + key: edge[key] for key in ["name", "source_label", "target_label", "properties"] if key in edge } mini_schema["edgelabels"].append(new_edge) 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 14064af18..8021a95b6 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 @@ -41,6 +41,7 @@ class BuildSemanticIndex: +<<<<<<< HEAD <<<<<<< HEAD def __init__(self, embedding: BaseEmbedding): self.folder_name = get_index_folder_name( @@ -55,6 +56,10 @@ def __init__(self, embedding: BaseEmbedding): 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") >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + 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") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) self.embedding = embedding self.sm = SchemaManager(huge_settings.graph_name) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index 28baca370..d76ad021a 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -43,7 +43,10 @@ def __init__( if not vector_index.exist("gremlin_examples"): log.warning("No gremlin example index found, will generate one.") self.vector_index = vector_index.from_name(self.embedding.get_embedding_dim(), "gremlin_examples") +<<<<<<< HEAD +======= +>>>>>>> 8e0bf08 (chore: mark vectordb optional) self._build_default_example_index() else: self.vector_index = vector_index.from_name(self.embedding.get_embedding_dim(), "gremlin_examples") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py index ab5a93912..8a2342c1a 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py @@ -40,6 +40,7 @@ def __init__(self, vector_index: type[VectorStoreBase], embedding: BaseEmbedding self.embedding = embedding self.topk = topk <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space @@ -52,6 +53,9 @@ def __init__(self, vector_index: type[VectorStoreBase], embedding: BaseEmbedding ======= self.vector_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "chunks") >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + self.vector_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "chunks") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query = context.get("query") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py index df697aad0..b5581b256 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py @@ -62,11 +62,13 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: +<<<<<<< HEAD context_str = f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") +======= + context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip("\n") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) - final_prompt = self._prompt_template.format( - context_str=context_str, query_str=self._question - ) + final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) response = self._llm.generate(prompt=final_prompt) return {"answer": response} @@ -96,9 +98,7 @@ def handle_vector_graph(self, context): vector_result_context = "No (vector)phrase related to the query." graph_result = context.get("graph_result") if graph_result: - graph_context_head = context.get( - "graph_context_head", "Knowledge from graphdb for the query:\n" - ) + graph_context_head = context.get("graph_context_head", "Knowledge from graphdb for the query:\n") graph_result_context = graph_context_head + "\n".join( f"{i + 1}. {res}" for i, res in enumerate(graph_result) ) @@ -111,11 +111,13 @@ async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[st context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: +<<<<<<< HEAD context_str = f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") +======= + context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip("\n") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) - final_prompt = self._prompt_template.format( - context_str=context_str, query_str=self._question - ) + final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) response = self._llm.generate(prompt=final_prompt) yield {"answer": response} return @@ -141,28 +143,32 @@ async def async_generate( final_prompt = self._question async_tasks["raw_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) if self._vector_only_answer: +<<<<<<< HEAD context_str = f"{context_head_str}\n" f"{vector_result_context}\n" f"{context_tail_str}".strip("\n") +======= + context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip("\n") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) - final_prompt = self._prompt_template.format( - context_str=context_str, query_str=self._question - ) - async_tasks["vector_only_task"] = asyncio.create_task( - self._llm.agenerate(prompt=final_prompt) - ) + final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + async_tasks["vector_only_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) if self._graph_only_answer: +<<<<<<< HEAD context_str = f"{context_head_str}\n" f"{graph_result_context}\n" f"{context_tail_str}".strip("\n") +======= + context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip("\n") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) - final_prompt = self._prompt_template.format( - context_str=context_str, query_str=self._question - ) - async_tasks["graph_only_task"] = asyncio.create_task( - self._llm.agenerate(prompt=final_prompt) - ) + final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + async_tasks["graph_only_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) if self._graph_vector_answer: context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" +<<<<<<< HEAD context_str = f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") +======= + context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) async_tasks["graph_vector_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) @@ -198,17 +204,17 @@ async def async_streaming_generate( if self._raw_answer: final_prompt = self._question async_generators.append( - self.__llm_generate_with_meta_info( - task_id=auto_id, target_key="raw_answer", prompt=final_prompt - ) + self.__llm_generate_with_meta_info(task_id=auto_id, target_key="raw_answer", prompt=final_prompt) ) auto_id += 1 if self._vector_only_answer: +<<<<<<< HEAD context_str = f"{context_head_str}\n" f"{vector_result_context}\n" f"{context_tail_str}".strip("\n") +======= + context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip("\n") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) - final_prompt = self._prompt_template.format( - context_str=context_str, query_str=self._question - ) + final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) async_generators.append( self.__llm_generate_with_meta_info( task_id=auto_id, target_key="vector_only_answer", prompt=final_prompt @@ -216,26 +222,28 @@ async def async_streaming_generate( ) auto_id += 1 if self._graph_only_answer: +<<<<<<< HEAD context_str = f"{context_head_str}\n" f"{graph_result_context}\n" f"{context_tail_str}".strip("\n") +======= + context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip("\n") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) - final_prompt = self._prompt_template.format( - context_str=context_str, query_str=self._question - ) + final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) async_generators.append( - self.__llm_generate_with_meta_info( - task_id=auto_id, target_key="graph_only_answer", prompt=final_prompt - ) + self.__llm_generate_with_meta_info(task_id=auto_id, target_key="graph_only_answer", prompt=final_prompt) ) auto_id += 1 if self._graph_vector_answer: context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" +<<<<<<< HEAD context_str = f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") +======= + context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) - final_prompt = self._prompt_template.format( - context_str=context_str, query_str=self._question - ) + final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) async_generators.append( self.__llm_generate_with_meta_info( task_id=auto_id, target_key="graph_vector_answer", prompt=final_prompt diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py index 981123954..82d985a38 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py @@ -51,6 +51,7 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: llm_output = self.llm.generate(prompt=prompt) data["triples"] = [] extract_triples_by_regex(llm_output, data) +<<<<<<< HEAD <<<<<<< HEAD print( f"LLM {self.__class__.__name__} input:{prompt} \n" @@ -59,6 +60,9 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: ======= print(f"LLM {self.__class__.__name__} input:{prompt} \n" f" output: {llm_output} \n data: {data}") >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + print(f"LLM {self.__class__.__name__} input:{prompt} \n output: {llm_output} \n data: {data}") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) data["call_count"] = data.get("call_count", 0) + 1 return data diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index edea66f26..e1d299cc1 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -53,6 +53,7 @@ def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional return None example_strings = [] for example in examples: +<<<<<<< HEAD <<<<<<< HEAD example_strings.append( f"- query: {example['query']}\n" @@ -61,6 +62,9 @@ def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional ======= example_strings.append(f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```") >>>>>>> 38dce0b (feat(llm): vector db finished) +======= + example_strings.append(f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) return "\n\n".join(example_strings) def _format_vertices(self, vertices: Optional[List[str]]) -> Optional[str]: @@ -94,9 +98,7 @@ async def async_generate(self, context: Dict[str, Any]): vertices=self._format_vertices(vertices=self.vertices), properties=self._format_properties(properties=None), ) - async_tasks["initialized_answer"] = asyncio.create_task( - self.llm.agenerate(prompt=init_prompt) - ) + async_tasks["initialized_answer"] = asyncio.create_task(self.llm.agenerate(prompt=init_prompt)) raw_response = await async_tasks["raw_answer"] initialized_response = await async_tasks["initialized_answer"] 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 504d683df..362670a4b 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 @@ -75,6 +75,7 @@ def generate_extract_triple_prompt(text, schema=None) -> str: if schema: return schema_real_prompt +<<<<<<< HEAD log.warning( <<<<<<< HEAD "Recommend to provide a graph schema to improve the extraction accuracy. " @@ -83,6 +84,9 @@ def generate_extract_triple_prompt(text, schema=None) -> str: "Recommend to provide a graph schema to improve the extraction accuracy. " "Now using the default schema." >>>>>>> 87ee5d3 (style: format code with black line-length 120) ) +======= + log.warning("Recommend to provide a graph schema to improve the extraction accuracy. Now using the default schema.") +>>>>>>> 8e0bf08 (chore: mark vectordb optional) return text_based_prompt @@ -111,6 +115,7 @@ def extract_triples_by_regex_with_schema(schema, text, graph): # TODO: use a more efficient way to compare the extract & input property p_lower = p.lower() for vertex in schema["vertices"]: +<<<<<<< HEAD <<<<<<< HEAD if vertex["vertex_label"] == label and any( pp.lower() == p_lower for pp in vertex["properties"] @@ -118,6 +123,9 @@ def extract_triples_by_regex_with_schema(schema, text, graph): ======= if vertex["vertex_label"] == label and any(pp.lower() == p_lower for pp in vertex["properties"]): >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + if vertex["vertex_label"] == label and any(pp.lower() == p_lower for pp in vertex["properties"]): +>>>>>>> 8e0bf08 (chore: mark vectordb optional) id = f"{label}-{s}" if id not in vertices_dict: vertices_dict[id] = { @@ -216,6 +224,7 @@ def valid(self, element_id: str, max_length: int = 256) -> bool: def _filter_long_id(self, graph) -> Dict[str, List[Any]]: graph["vertices"] = [vertex for vertex in graph["vertices"] if self.valid(vertex["id"])] +<<<<<<< HEAD <<<<<<< HEAD graph["edges"] = [ edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"]) @@ -223,4 +232,7 @@ def _filter_long_id(self, graph) -> Dict[str, List[Any]]: ======= graph["edges"] = [edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"])] >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + graph["edges"] = [edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"])] +>>>>>>> 8e0bf08 (chore: mark vectordb optional) return graph diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index 271b530c6..0aa2a2bea 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -85,6 +85,7 @@ def _extract_keywords_from_response( for match in matches: match = match[len(start_token) :].strip() +<<<<<<< HEAD <<<<<<< HEAD keywords.extend( k.lower() if lowercase else k @@ -94,13 +95,14 @@ def _extract_keywords_from_response( ======= keywords.extend(k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1) >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + keywords.extend(k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1) +>>>>>>> 8e0bf08 (chore: mark vectordb optional) # if the keyword consists of multiple words, split into sub-words (removing stopwords) results = set(keywords) for token in keywords: sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: - results.update( - w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language) - ) + results.update(w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language)) return results 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 0b34717e3..710760c7d 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 @@ -125,12 +125,16 @@ 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( +<<<<<<< HEAD <<<<<<< HEAD "Invalid property graph! No JSON object found, " "please check the output format example in prompt." ======= "Invalid property graph! No JSON object found, " "please check the output format example in prompt." >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + "Invalid property graph! No JSON object found, please check the output format example in prompt." +>>>>>>> 8e0bf08 (chore: mark vectordb optional) ) return [] json_str = json_match.group(1).strip() @@ -139,11 +143,7 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: try: property_graph = json.loads(json_str) # 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 - ): + 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'.") return items @@ -171,7 +171,5 @@ def process_items(item_list, valid_labels, item_type): process_items(property_graph["vertices"], vertex_label_set, "vertex") process_items(property_graph["edges"], edge_label_set, "edge") except json.JSONDecodeError: - log.critical( - "Invalid property graph JSON! Please check the extracted JSON data carefully" - ) + log.critical("Invalid property graph JSON! Please check the extracted JSON data carefully") return items diff --git a/hugegraph-llm/src/hugegraph_llm/utils/anchor.py b/hugegraph-llm/src/hugegraph_llm/utils/anchor.py index 4542a7fd9..a46a4e499 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/anchor.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/anchor.py @@ -35,6 +35,5 @@ def get_project_root() -> Path: return parent # Raise an error if no project root is found raise RuntimeError( - "Project root could not be determined. " - "Ensure that 'pyproject.toml' or '.git' exists in the project directory." + "Project root could not be determined. Ensure that 'pyproject.toml' or '.git' exists in the project directory." ) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py index b2f485cea..45eb18626 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py @@ -24,9 +24,7 @@ from hugegraph_llm.models.embeddings.base import BaseEmbedding -async def _get_batch_with_progress( - embedding: BaseEmbedding, batch: list[str], pbar: tqdm -) -> list[Any]: +async def _get_batch_with_progress(embedding: BaseEmbedding, batch: list[str], pbar: tqdm) -> list[Any]: result = await embedding.async_get_texts_embeddings(batch) pbar.update(1) return result diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index 28b738b13..bf3b84ecf 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -53,9 +53,7 @@ def init_hg_test_data(): schema = client.schema() schema.propertyKey("name").asText().ifNotExist().create() schema.propertyKey("birthDate").asText().ifNotExist().create() - schema.vertexLabel("Person").properties( - "name", "birthDate" - ).useCustomizeStringId().ifNotExist().create() + schema.vertexLabel("Person").properties("name", "birthDate").useCustomizeStringId().ifNotExist().create() schema.vertexLabel("Movie").properties("name").useCustomizeStringId().ifNotExist().create() schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() @@ -140,6 +138,7 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): elif filename == "vertices.json": data_full = client.gremlin().exec(query)["data"][0]["vertices"] data = ( +<<<<<<< HEAD <<<<<<< HEAD [ {key: value for key, value in vertex.items() if key != "id"} @@ -148,6 +147,9 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): ======= [{key: value for key, value in vertex.items() if key != "id"} for vertex in data_full] >>>>>>> 87ee5d3 (style: format code with black line-length 120) +======= + [{key: value for key, value in vertex.items() if key != "id"} for vertex in data_full] +>>>>>>> 8e0bf08 (chore: mark vectordb optional) if all_pk_flag else data_full ) @@ -156,9 +158,7 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): data_full = query if isinstance(data_full, dict) and "schema" in data_full: groovy_filename = filename.replace(".json", ".groovy") - with open( - os.path.join(backup_subdir, groovy_filename), "w", encoding="utf-8" - ) as groovy_file: + with open(os.path.join(backup_subdir, groovy_filename), "w", encoding="utf-8") as groovy_file: groovy_file.write(str(data_full["schema"])) else: data = data_full From 6d088f7be49bcf1f044f2e2469561cb7c97ca4b9 Mon Sep 17 00:00:00 2001 From: imbajin Date: Wed, 27 Aug 2025 21:47:41 +0800 Subject: [PATCH 23/71] fix cycle import & add docs --- hugegraph-llm/config.md | 13 +++++ hugegraph-llm/pyproject.toml | 9 ++++ .../src/hugegraph_llm/config/__init__.py | 3 +- .../src/hugegraph_llm/config/index_config.py | 2 +- .../demo/rag_demo/configs_block.py | 50 ++++++++++++++----- .../hugegraph_llm/demo/rag_demo/rag_block.py | 8 +-- .../demo/rag_demo/text2gremlin_block.py | 12 +++++ .../hugegraph_llm/utils/graph_index_utils.py | 13 ++--- .../hugegraph_llm/utils/vector_index_utils.py | 13 +++-- 9 files changed, 92 insertions(+), 31 deletions(-) diff --git a/hugegraph-llm/config.md b/hugegraph-llm/config.md index a55172f33..5b0e766d5 100644 --- a/hugegraph-llm/config.md +++ b/hugegraph-llm/config.md @@ -16,6 +16,7 @@ - [LiteLLM 配置](#litellm-配置) - [重排序配置](#重排序配置) - [HugeGraph 数据库配置](#hugegraph-数据库配置) + - [向量数据库配置](#向量数据库配置) - [管理员配置](#管理员配置) - [配置使用示例](#配置使用示例) - [配置文件位置](#配置文件位置) @@ -127,6 +128,18 @@ | `TOPK_PER_KEYWORD` | Optional[Integer] | 1 | 每个关键词返回的 TopK 数量 | | `TOPK_RETURN_RESULTS` | Optional[Integer] | 20 | 返回结果数量 | +### 向量数据库配置 + +| 配置项 | 类型 | 默认值 | 说明 | +|------------------|------------------|-------|------------------------| +| `QDRANT_HOST` | Optional[String] | None | Qdrant 服务器主机地址 | +| `QDRANT_PORT` | Integer | 6333 | Qdrant 服务器端口 | +| `QDRANT_API_KEY` | Optional[String] | None | Qdrant API 密钥（如果设置了的话） | +| `MILVUS_HOST` | Optional[String] | None | Milvus 服务器主机地址 | +| `MILVUS_PORT` | Integer | 19530 | Milvus 服务器端口 | +| `MILVUS_USER` | String | "" | Milvus 用户名 | +| `MILVUS_PASSWORD`| String | "" | Milvus 密码 | + ### 管理员配置 | 配置项 | 类型 | 默认值 | 说明 | diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index 0d12210fb..224cf0306 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -67,6 +67,15 @@ vectordb = [ "pymilvus==2.5.9", "qdrant-client==1.14.2", ] +======= + + # Vector database dependencies + "pymilvus==2.5.9", + "qdrant-client==1.14.2", + + ] + +>>>>>>> a255aed (fix cycle import & add docs) [project.urls] homepage = "https://hugegraph.apache.org/" repository = "https://github.com/apache/incubator-hugegraph-ai" diff --git a/hugegraph-llm/src/hugegraph_llm/config/__init__.py b/hugegraph-llm/src/hugegraph_llm/config/__init__.py index f7f9cf290..43efb0ab3 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/__init__.py +++ b/hugegraph-llm/src/hugegraph_llm/config/__init__.py @@ -20,10 +20,9 @@ import os -from hugegraph_llm.config.index_config import IndexConfig - from .admin_config import AdminConfig from .hugegraph_config import HugeGraphConfig +from .index_config import IndexConfig from .llm_config import LLMConfig from .prompt_config import PromptConfig diff --git a/hugegraph-llm/src/hugegraph_llm/config/index_config.py b/hugegraph-llm/src/hugegraph_llm/config/index_config.py index afe84a793..b0a6e21e1 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/index_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/index_config.py @@ -33,4 +33,4 @@ class IndexConfig(BaseConfig): milvus_user: str = os.environ.get("MILVUS_USER", "") milvus_password: str = os.environ.get("MILVUS_PASSWORD", "") - now_vector_index: str = "Faiss" + cur_vector_index: str = "Faiss" diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index d43bee8b8..464b339ca 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -71,10 +71,7 @@ def _is_private_ipv4(host: str) -> bool: def test_litellm_embedding(api_key, api_base, model_name, model_dim) -> int: llm_client = LiteLLMEmbedding( -<<<<<<< HEAD -======= embedding_dimension=model_dim, ->>>>>>> 38dce0b (feat(llm): vector db finished) api_key=api_key, api_base=api_base, model_name=model_name, @@ -354,13 +351,16 @@ def create_configs_block() -> list: with gr.Row(): graph_config_input = [ gr.Textbox( - value=lambda: huge_settings.graph_url, + value=huge_settings.graph_url, label="url", info="IP:PORT (e.g. 127.0.0.1:8080) or full URL (e.g. http://127.0.0.1:8080)", ), gr.Textbox( <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> a255aed (fix cycle import & add docs) value=huge_settings.graph_name, label="graph", info="The graph name of HugeGraph-Server instance", @@ -372,6 +372,7 @@ def create_configs_block() -> list: ), gr.Textbox( value=huge_settings.graph_pwd, +<<<<<<< HEAD label="pwd", type="password", info="Password for graph server auth", @@ -390,13 +391,15 @@ def create_configs_block() -> list: >>>>>>> 38dce0b (feat(llm): vector db finished) ======= value=lambda: huge_settings.graph_pwd, +======= +>>>>>>> a255aed (fix cycle import & add docs) label="pwd", type="password", info="Password for graph server auth", >>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=lambda: huge_settings.graph_space, + value=huge_settings.graph_space, label="graphspace (Optional)", info="Namespace for multi-tenant scenarios (leave empty if not using graphspaces)", ), @@ -432,6 +435,9 @@ def chat_llm_settings(llm_type): gr.Textbox( <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> a255aed (fix cycle import & add docs) value=getattr(llm_settings, "openai_chat_api_key"), label="api_key", type="password", @@ -510,6 +516,7 @@ def chat_llm_settings(llm_type): llm_config_input = [ gr.Textbox( <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "litellm_chat_api_key"), label="api_key", @@ -519,17 +526,20 @@ def chat_llm_settings(llm_type): >>>>>>> 38dce0b (feat(llm): vector db finished) ======= value=lambda: getattr(llm_settings, "litellm_chat_api_key"), +======= + value=getattr(llm_settings, "litellm_chat_api_key"), +>>>>>>> a255aed (fix cycle import & add docs) label="api_key", type="password", >>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_chat_api_base"), + value=getattr(llm_settings, "litellm_chat_api_base"), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_chat_language_model"), + value=getattr(llm_settings, "litellm_chat_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -590,6 +600,9 @@ def extract_llm_settings(llm_type): gr.Textbox( <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> a255aed (fix cycle import & add docs) value=getattr(llm_settings, "openai_extract_api_key"), label="api_key", type="password", @@ -670,6 +683,7 @@ def extract_llm_settings(llm_type): llm_config_input = [ gr.Textbox( <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "litellm_extract_api_key"), label="api_key", @@ -679,17 +693,20 @@ def extract_llm_settings(llm_type): >>>>>>> 38dce0b (feat(llm): vector db finished) ======= value=lambda: getattr(llm_settings, "litellm_extract_api_key"), +======= + value=getattr(llm_settings, "litellm_extract_api_key"), +>>>>>>> a255aed (fix cycle import & add docs) label="api_key", type="password", >>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_extract_api_base"), + value=getattr(llm_settings, "litellm_extract_api_base"), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_extract_language_model"), + value=getattr(llm_settings, "litellm_extract_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -736,6 +753,9 @@ def text2gql_llm_settings(llm_type): gr.Textbox( <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> a255aed (fix cycle import & add docs) value=getattr(llm_settings, "openai_text2gql_api_key"), label="api_key", type="password", @@ -1057,16 +1077,20 @@ def reranker_settings(reranker_type): inputs=reranker_config_input, # pylint: disable=no-member ) +<<<<<<< HEAD <<<<<<< HEAD ======= with gr.Accordion("5. Set up the vector database.", open=False): +======= + with gr.Accordion("5. Set up the vector engine.", open=False): +>>>>>>> a255aed (fix cycle import & add docs) engine_selector = gr.Dropdown( choices=["Faiss", "Milvus", "Qdrant"], - value=lambda: index_settings.now_vector_index, - label="Select vector database.", + value=index_settings.cur_vector_index, + label="Select vector engine.", ) engine_selector.select( - fn=lambda engine: setattr(index_settings, "now_vector_index", engine), + fn=lambda engine: setattr(index_settings, "cur_vector_index", engine), inputs=[engine_selector], ) >>>>>>> 38dce0b (feat(llm): vector db finished) @@ -1078,7 +1102,7 @@ def get_header_with_language_indicator(language: str) -> str: language_class = language.lower() if language == "CN": - title_text = "当前prompt语言: 中文 (CN)" + title_text = "当前 prompt 语言：中文 (CN)" else: title_text = "Current prompt Language: English (EN)" html_content = f""" diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 396727f82..78cc7e242 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -82,10 +82,10 @@ def rag_answer( ======= rag = RAGPipeline() if vector_search: - rag.query_vector_index(vector_index_str=index_settings.now_vector_index) + rag.query_vector_index(vector_index_str=index_settings.cur_vector_index) if graph_search: rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid( - vector_index_str=index_settings.now_vector_index, + vector_index_str=index_settings.cur_vector_index, vector_dis_threshold=vector_dis_threshold, topk_per_keyword=topk_per_keyword, ).import_schema(huge_settings.graph_name).query_graphdb( @@ -228,10 +228,10 @@ async def rag_answer_streaming( ======= rag = RAGPipeline() if vector_search: - rag.query_vector_index(vector_index_str=index_settings.now_vector_index) + rag.query_vector_index(vector_index_str=index_settings.cur_vector_index) if graph_search: rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid( - vector_index_str=index_settings.now_vector_index + vector_index_str=index_settings.cur_vector_index ).import_schema(huge_settings.graph_name).query_graphdb( num_gremlin_generate_example=gremlin_tmpl_num, gremlin_prompt=gremlin_prompt, diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index fe885bfa9..f3119f67c 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -98,6 +98,7 @@ def store_schema(schema, question, gremlin_prompt): def build_example_vector_index(temp_file) -> dict: +<<<<<<< HEAD <<<<<<< HEAD folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) index_path = os.path.join(resource_path, folder_name, "gremlin_examples") @@ -109,6 +110,9 @@ def build_example_vector_index(temp_file) -> dict: assert vector_index, 'vector db name is error' >>>>>>> 38dce0b (feat(llm): vector db finished) ======= +======= + vector_index = get_vector_index_class(index_settings.cur_vector_index) +>>>>>>> a255aed (fix cycle import & add docs) assert vector_index, "vector db name is error" >>>>>>> 87ee5d3 (style: format code with black line-length 120) if temp_file is None: @@ -219,12 +223,16 @@ def _execute_queries(context, output_types): def gremlin_generate( inp, example_num, schema, gremlin_prompt, requested_outputs: Optional[List[str]] = None ) -> GremlinResult: +<<<<<<< HEAD <<<<<<< HEAD generator = GremlinGenerator( llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() ) ======= vector_index = get_vector_index_class(index_settings.now_vector_index) +======= + vector_index = get_vector_index_class(index_settings.cur_vector_index) +>>>>>>> a255aed (fix cycle import & add docs) generator = GremlinGenerator(llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding()) >>>>>>> 8e0bf08 (chore: mark vectordb optional) sm = SchemaManager(graph_name=schema) @@ -406,12 +414,16 @@ def graph_rag_recall( rag.extract_keywords().keywords_to_vid( <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD ======= vector_index=index_settings.now_vector_index, >>>>>>> 38dce0b (feat(llm): vector db finished) ======= vector_index_str=index_settings.now_vector_index, >>>>>>> dd3b085 (feat(llm): nexpected-keyword-arg,unused-import) +======= + vector_index_str=index_settings.cur_vector_index, +>>>>>>> a255aed (fix cycle import & add docs) vector_dis_threshold=vector_dis_threshold, topk_per_keyword=topk_per_keyword, ) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index ddb531508..8a2130218 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -50,6 +50,7 @@ def get_graph_index_info(): builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) graph_summary_info = builder.fetch_graph_data().run() +<<<<<<< HEAD <<<<<<< HEAD folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( @@ -59,7 +60,10 @@ def get_graph_index_info(): str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix, record_miss=False ) ======= - vector_index = get_vector_index_class(index_settings.now_vector_index) + vector_index = get_vector_index_class(index_settings.cur_vector_index) +======= + vector_index = get_vector_index_class(index_settings.cur_vector_index) +>>>>>>> a255aed (fix cycle import & add docs) vector_index_entity = vector_index.from_name( Embeddings().get_embedding().get_embedding_dim(), huge_settings.graph_name, "chunks" ) @@ -172,14 +176,7 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: def update_vid_embedding(): -<<<<<<< HEAD scheduler = SchedulerSingleton.get_instance() -======= - vector_index = get_vector_index_class(index_settings.now_vector_index) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - builder.fetch_graph_data().build_vertex_id_semantic_index(vector_index) - log.debug("Operators: %s", builder.operators) ->>>>>>> 38dce0b (feat(llm): vector db finished) try: return scheduler.schedule_flow("update_vid_embeddings") except Exception as e: # pylint: disable=broad-exception-caught diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 4a1363676..affcf2459 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -72,26 +72,33 @@ def read_documents(input_file, input_text): # pylint: disable=C0301 def get_vector_index_info(): - vector_index = get_vector_index_class(index_settings.now_vector_index) + vector_index = get_vector_index_class(index_settings.cur_vector_index) vector_index_entity = vector_index.from_name( Embeddings().get_embedding().get_embedding_dim(), huge_settings.graph_name, "chunks" ) return json.dumps( +<<<<<<< HEAD {**vector_index_entity.get_vector_index_info(), "now_vector_index": index_settings.now_vector_index}, +======= + { + **vector_index_entity.get_vector_index_info(), + "cur_vector_index": index_settings.cur_vector_index, + }, +>>>>>>> a255aed (fix cycle import & add docs) ensure_ascii=False, indent=2, ) def clean_vector_index(): - vector_index = get_vector_index_class(index_settings.now_vector_index) + vector_index = get_vector_index_class(index_settings.cur_vector_index) vector_index.clean(huge_settings.graph_name, "chunks") gr.Info("Clean vector index successfully!") def build_vector_index(input_file, input_text): - vector_index = get_vector_index_class(index_settings.now_vector_index) + vector_index = get_vector_index_class(index_settings.cur_vector_index) if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) From 12bf4151e7d119a607aceab4c84ca31aae1c6a54 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Thu, 28 Aug 2025 17:26:37 +0800 Subject: [PATCH 24/71] fix --- hugegraph-llm/pyproject.toml | 1 + .../hugegraph_llm/config/models/__init__.py | 2 +- .../demo/rag_demo/admin_block.py | 3 +- .../demo/rag_demo/configs_block.py | 37 ------------------- .../hugegraph_llm/demo/rag_demo/rag_block.py | 5 +-- 5 files changed, 5 insertions(+), 43 deletions(-) diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index 224cf0306..81e3dc1f2 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -40,6 +40,7 @@ dependencies = [ "numpy", "pandas", "pydantic", + "tqdm", # LLM specific dependencies "openai", diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py b/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py index d7738036e..087d89477 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py @@ -18,4 +18,4 @@ from .base_config import BaseConfig from .base_prompt_config import BasePromptConfig -__all__ = ["BaseConfig", "BasePromptConfig"] \ No newline at end of file +__all__ = ["BaseConfig", "BasePromptConfig"] diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 78ac89ad9..0b041387e 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -19,7 +19,6 @@ from collections import deque import gradio as gr -from gradio import Request from hugegraph_llm.config import admin_settings from hugegraph_llm.utils.log import log @@ -70,7 +69,7 @@ def clear_llm_server_log(): # Function to validate password and control access to logs -def check_password(password, request: Request = None): +def check_password(password, request = None): client_ip = request.client.host if request else "Unknown IP" admin_token = admin_settings.admin_token diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 464b339ca..8d90da414 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -17,7 +17,6 @@ import json import os -import urllib.parse as _urlparse from functools import partial from typing import Optional @@ -34,41 +33,6 @@ current_llm = "chat" -def _validate_url_safe(url: str) -> None: - """Basic SSRF guard: allow only http/https and forbid unexpected schemes.""" - parsed = _urlparse.urlparse(url) - if parsed.scheme not in {"http", "https"}: - raise gr.Error("Only http/https URLs are allowed for connection test.") - if not parsed.netloc: - raise gr.Error("URL missing hostname.") - - # 防止 SSRF：禁止访问本地主机或内网地址 - hostname = parsed.hostname or "" - # IPv4 私有网段 - private_ipv4_networks = [ - ("10.",), - ("172.", range(16, 32)), - ("192.168.",), - ("127.",), - ("0.",), - ] - - def _is_private_ipv4(host: str) -> bool: - for prefix in private_ipv4_networks: - base = prefix[0] - if host.startswith(base): - # 处理 172.16.0.0/12 特例 - if len(prefix) == 1: - return True - if int(host.split(".")[1]) in prefix[1]: - return True - return False - - # IPv6 localhost - if hostname in {"localhost", "::1"} or _is_private_ipv4(hostname): - raise gr.Error("Connection to localhost or private network addresses is not allowed.") - - def test_litellm_embedding(api_key, api_base, model_name, model_dim) -> int: llm_client = LiteLLMEmbedding( embedding_dimension=model_dim, @@ -104,7 +68,6 @@ def test_litellm_chat(api_key, api_base, model_name, max_tokens: int) -> int: def test_api_connection(url, method="GET", headers=None, params=None, body=None, auth=None, origin_call=None) -> int: # TODO: use fastapi.request / starlette instead? log.debug("Request URL: %s", url) - _validate_url_safe(url) try: if method.upper() == "GET": resp = requests.get(url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 78cc7e242..03720d810 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -18,12 +18,11 @@ # pylint: disable=E1101 import os -from typing import AsyncGenerator, Literal, Optional, Tuple +from typing import Any, AsyncGenerator, Literal, Optional, Tuple import gradio as gr from hugegraph_llm.flows.scheduler import SchedulerSingleton import pandas as pd -from gradio.utils import NamedString <<<<<<< HEAD from hugegraph_llm.config import resource_path, prompt, llm_settings @@ -437,7 +436,7 @@ def toggle_slider(enable): resource_path, "demo", "questions_template.xlsx" ) - def read_file_to_excel(file: NamedString, line_count: Optional[int] = None): + def read_file_to_excel(file: Any, line_count: Optional[int] = None): df = None if not file: return pd.DataFrame(), 1 From 6d7c9ed90d4ac8f9149134b7040b4d88f6f4658d Mon Sep 17 00:00:00 2001 From: lingxiao Date: Thu, 28 Aug 2025 17:34:32 +0800 Subject: [PATCH 25/71] fix --- hugegraph-llm/pyproject.toml | 1 + .../hugegraph_llm/api/models/rag_requests.py | 84 ++++++--- .../src/hugegraph_llm/api/rag_api.py | 31 +++- .../src/hugegraph_llm/config/index_config.py | 4 +- .../src/hugegraph_llm/config/llm_config.py | 20 ++- .../config/models/base_config.py | 12 +- .../config/models/base_prompt_config.py | 12 +- .../demo/rag_demo/admin_block.py | 12 +- .../src/hugegraph_llm/demo/rag_demo/app.py | 12 +- .../demo/rag_demo/configs_block.py | 170 +++++++++++++++--- .../demo/rag_demo/other_block.py | 12 +- .../hugegraph_llm/demo/rag_demo/rag_block.py | 46 ++++- .../demo/rag_demo/text2gremlin_block.py | 59 ++++-- .../demo/rag_demo/vector_graph_block.py | 50 ++++-- .../src/hugegraph_llm/document/chunk_split.py | 10 ++ .../indices/vector_index/base.py | 4 +- .../vector_index/faiss_vector_store.py | 4 +- .../vector_index/milvus_vector_store.py | 16 +- .../vector_index/qdrant_vector_store.py | 16 +- .../hugegraph_llm/middleware/middleware.py | 4 +- .../models/embeddings/init_embedding.py | 8 +- .../hugegraph_llm/models/embeddings/openai.py | 4 +- .../src/hugegraph_llm/models/llms/init_llm.py | 6 +- .../src/hugegraph_llm/models/llms/ollama.py | 14 +- .../src/hugegraph_llm/models/llms/openai.py | 12 +- .../hugegraph_llm/models/rerankers/cohere.py | 12 +- .../models/rerankers/init_reranker.py | 4 +- .../models/rerankers/siliconflow.py | 12 +- .../operators/common_op/check_schema.py | 62 +++++-- .../operators/common_op/merge_dedup_rerank.py | 17 +- .../operators/common_op/nltk_helper.py | 4 +- .../operators/document_op/chunk_split.py | 4 +- .../operators/document_op/word_extract.py | 8 +- .../hugegraph_llm/operators/graph_rag_task.py | 10 +- .../operators/gremlin_generate_task.py | 18 +- .../hugegraph_op/commit_to_hugegraph.py | 78 +++++--- .../hugegraph_op/fetch_graph_data.py | 4 +- .../operators/hugegraph_op/graph_rag_query.py | 86 ++++++--- .../operators/hugegraph_op/schema_manager.py | 10 +- .../index_op/build_gremlin_example_index.py | 8 +- .../index_op/build_semantic_index.py | 35 +++- .../index_op/gremlin_example_index_query.py | 22 ++- .../operators/index_op/semantic_id_query.py | 10 +- .../operators/index_op/vector_index_query.py | 6 + .../operators/kg_construction_task.py | 4 +- .../operators/llm_op/answer_synthesize.py | 140 ++++++++++++--- .../operators/llm_op/disambiguate_data.py | 6 + .../operators/llm_op/gremlin_generate.py | 26 ++- .../operators/llm_op/info_extract.py | 23 +++ .../operators/llm_op/keyword_extract.py | 16 +- .../operators/llm_op/prompt_generate.py | 12 +- .../llm_op/property_graph_extract.py | 28 ++- .../operators/llm_op/schema_build.py | 8 +- .../llm_op/unstructured_data_utils.py | 8 +- .../src/hugegraph_llm/utils/decorators.py | 4 +- .../hugegraph_llm/utils/embedding_utils.py | 16 +- .../hugegraph_llm/utils/graph_index_utils.py | 83 ++++----- .../hugegraph_llm/utils/hugegraph_utils.py | 60 +++++-- .../hugegraph_llm/utils/vector_index_utils.py | 29 ++- 59 files changed, 1150 insertions(+), 346 deletions(-) diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index 81e3dc1f2..7724f7444 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -41,6 +41,7 @@ dependencies = [ "pandas", "pydantic", "tqdm", + "tqdm", # LLM specific dependencies "openai", diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py index 5222f0cfa..9ebc94eb1 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py +++ b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py @@ -34,16 +34,34 @@ class GraphConfigRequest(BaseModel): class RAGRequest(BaseModel): query: str = Query(..., description="Query you want to ask") raw_answer: bool = Query(False, description="Use LLM to generate answer directly") - vector_only: bool = Query(False, description="Use LLM to generate answer with vector") - graph_only: bool = Query(True, description="Use LLM to generate answer with graph RAG only") - graph_vector_answer: bool = Query(False, description="Use LLM to generate answer with vector & GraphRAG") - graph_ratio: float = Query(0.5, description="The ratio of GraphRAG ans & vector ans") - rerank_method: Literal["bleu", "reranker"] = Query("bleu", description="Method to rerank the results.") - near_neighbor_first: bool = Query(False, description="Prioritize near neighbors in the search results.") - custom_priority_info: str = Query("", description="Custom information to prioritize certain results.") + vector_only: bool = Query( + False, description="Use LLM to generate answer with vector" + ) + graph_only: bool = Query( + True, description="Use LLM to generate answer with graph RAG only" + ) + graph_vector_answer: bool = Query( + False, description="Use LLM to generate answer with vector & GraphRAG" + ) + graph_ratio: float = Query( + 0.5, description="The ratio of GraphRAG ans & vector ans" + ) + rerank_method: Literal["bleu", "reranker"] = Query( + "bleu", description="Method to rerank the results." + ) + near_neighbor_first: bool = Query( + False, description="Prioritize near neighbors in the search results." + ) + custom_priority_info: str = Query( + "", description="Custom information to prioritize certain results." + ) # Graph Configs - max_graph_items: int = Query(30, description="Maximum number of items for GQL queries in graph.") - topk_return_results: int = Query(20, description="Number of sorted results to return finally.") + max_graph_items: int = Query( + 30, description="Maximum number of items for GQL queries in graph." + ) + topk_return_results: int = Query( + 20, description="Number of sorted results to return finally." + ) vector_dis_threshold: float = Query( 0.9, description="Threshold for vector similarity\ @@ -54,10 +72,14 @@ class RAGRequest(BaseModel): description="TopK results returned for each keyword \ extracted from the query, by default only the most similar one is returned.", ) - client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") + client_config: Optional[GraphConfigRequest] = Query( + None, description="hugegraph server config." + ) # Keep prompt params in the end - answer_prompt: Optional[str] = Query(prompt.answer_prompt, description="Prompt to guide the answer generation.") + answer_prompt: Optional[str] = Query( + prompt.answer_prompt, description="Prompt to guide the answer generation." + ) keywords_extract_prompt: Optional[str] = Query( prompt.keywords_extract_prompt, description="Prompt for extracting keywords from query.", @@ -73,8 +95,12 @@ class RAGRequest(BaseModel): class GraphRAGRequest(BaseModel): query: str = Query(..., description="Query you want to ask") # Graph Configs - max_graph_items: int = Query(30, description="Maximum number of items for GQL queries in graph.") - topk_return_results: int = Query(20, description="Number of sorted results to return finally.") + max_graph_items: int = Query( + 30, description="Maximum number of items for GQL queries in graph." + ) + topk_return_results: int = Query( + 20, description="Number of sorted results to return finally." + ) vector_dis_threshold: float = Query( 0.9, description="Threshold for vector similarity \ @@ -86,16 +112,26 @@ class GraphRAGRequest(BaseModel): from the query, by default only the most similar one is returned.", ) - client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") - get_vertex_only: bool = Query(False, description="return only keywords & vertex (early stop).") + client_config: Optional[GraphConfigRequest] = Query( + None, description="hugegraph server config." + ) + get_vertex_only: bool = Query( + False, description="return only keywords & vertex (early stop)." + ) gremlin_tmpl_num: int = Query( 1, description="Number of Gremlin templates to use. If num <=0 means template is not provided", ) - rerank_method: Literal["bleu", "reranker"] = Query("bleu", description="Method to rerank the results.") - near_neighbor_first: bool = Query(False, description="Prioritize near neighbors in the search results.") - custom_priority_info: str = Query("", description="Custom information to prioritize certain results.") + rerank_method: Literal["bleu", "reranker"] = Query( + "bleu", description="Method to rerank the results." + ) + near_neighbor_first: bool = Query( + False, description="Prioritize near neighbors in the search results." + ) + custom_priority_info: str = Query( + "", description="Custom information to prioritize certain results." + ) gremlin_prompt: Optional[str] = Query( prompt.gremlin_generate_prompt, description="Prompt for the Text2Gremlin query.", @@ -139,12 +175,16 @@ class GremlinOutputType(str, Enum): class GremlinGenerateRequest(BaseModel): query: str - example_num: Optional[int] = Query(0, description="Number of Gremlin templates to use.(0 means no templates)") + example_num: Optional[int] = Query( + 0, description="Number of Gremlin templates to use.(0 means no templates)" + ) gremlin_prompt: Optional[str] = Query( prompt.gremlin_generate_prompt, description="Prompt for the Text2Gremlin query.", ) - client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") + client_config: Optional[GraphConfigRequest] = Query( + None, description="hugegraph server config." + ) output_types: Optional[List[GremlinOutputType]] = Query( default=[GremlinOutputType.TEMPLATE_GREMLIN], description=""" @@ -161,5 +201,7 @@ def validate_prompt_placeholders(cls, v): required_placeholders = ["{query}", "{schema}", "{example}", "{vertices}"] missing = [p for p in required_placeholders if p not in v] if missing: - raise ValueError(f"Prompt template is missing required placeholders: {', '.join(missing)}") + raise ValueError( + f"Prompt template is missing required placeholders: {', '.join(missing)}" + ) return v diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py index bcd15f54f..1d5b451b1 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py @@ -67,7 +67,8 @@ def rag_answer_api(req: RAGRequest): # Keep prompt params in the end custom_related_information=req.custom_priority_info, answer_prompt=req.answer_prompt or prompt.answer_prompt, - keywords_extract_prompt=req.keywords_extract_prompt or prompt.keywords_extract_prompt, + keywords_extract_prompt=req.keywords_extract_prompt + or prompt.keywords_extract_prompt, gremlin_prompt=req.gremlin_prompt or prompt.gremlin_generate_prompt, ) # TODO: we need more info in the response for users to understand the query logic @@ -135,7 +136,9 @@ def graph_rag_recall_api(req: GraphRAGRequest): except TypeError as e: log.error("TypeError in graph_rag_recall_api: %s", e) - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) + ) from e except Exception as e: log.error("Unexpected error occurred: %s", e) raise HTTPException( @@ -146,7 +149,9 @@ def graph_rag_recall_api(req: GraphRAGRequest): @router.post("/config/graph", status_code=status.HTTP_201_CREATED) def graph_config_api(req: GraphConfigRequest): # Accept status code - res = apply_graph_conf(req.url, req.name, req.user, req.pwd, req.gs, origin_call="http") + res = apply_graph_conf( + req.url, req.name, req.user, req.pwd, req.gs, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) # TODO: restructure the implement of llm to three types, like "/config/chat_llm" @@ -159,7 +164,9 @@ def llm_config_api(req: LLMConfigRequest): req.api_key, req.api_base, req.language_model, req.max_tokens, origin_call="http" ) else: - res = apply_llm_conf(req.host, req.port, req.language_model, None, origin_call="http") + res = apply_llm_conf( + req.host, req.port, req.language_model, None, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/embedding", status_code=status.HTTP_201_CREATED) @@ -167,9 +174,13 @@ def embedding_config_api(req: LLMConfigRequest): llm_settings.embedding_type = req.llm_type if req.llm_type == "openai": - res = apply_embedding_conf(req.api_key, req.api_base, req.language_model, origin_call="http") + res = apply_embedding_conf( + req.api_key, req.api_base, req.language_model, origin_call="http" + ) else: - res = apply_embedding_conf(req.host, req.port, req.language_model, origin_call="http") + res = apply_embedding_conf( + req.host, req.port, req.language_model, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/rerank", status_code=status.HTTP_201_CREATED) @@ -177,9 +188,13 @@ def rerank_config_api(req: RerankerConfigRequest): llm_settings.reranker_type = req.reranker_type if req.reranker_type == "cohere": - res = apply_reranker_conf(req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http") + res = apply_reranker_conf( + req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http" + ) elif req.reranker_type == "siliconflow": - res = apply_reranker_conf(req.api_key, req.reranker_model, None, origin_call="http") + res = apply_reranker_conf( + req.api_key, req.reranker_model, None, origin_call="http" + ) else: res = status.HTTP_501_NOT_IMPLEMENTED return generate_response(RAGResponse(status_code=res, message="Missing Value")) diff --git a/hugegraph-llm/src/hugegraph_llm/config/index_config.py b/hugegraph-llm/src/hugegraph_llm/config/index_config.py index b0a6e21e1..ad0db5975 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/index_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/index_config.py @@ -26,7 +26,9 @@ class IndexConfig(BaseConfig): qdrant_host: Optional[str] = os.environ.get("QDRANT_HOST", None) qdrant_port: int = int(os.environ.get("QDRANT_PORT", "6333")) - qdrant_api_key: Optional[str] = os.environ.get("QDRANT_API_KEY") if os.environ.get("QDRANT_API_KEY") else None + qdrant_api_key: Optional[str] = ( + os.environ.get("QDRANT_API_KEY") if os.environ.get("QDRANT_API_KEY") else None + ) milvus_host: Optional[str] = os.environ.get("MILVUS_HOST", None) milvus_port: int = int(os.environ.get("MILVUS_PORT", "19530")) diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index a63d401c1..8b4f274ee 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -32,16 +32,24 @@ class LLMConfig(BaseConfig): embedding_type: Optional[Literal["openai", "litellm", "ollama/local"]] = "openai" reranker_type: Optional[Literal["cohere", "siliconflow"]] = None # 1. OpenAI settings - openai_chat_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_chat_api_base: str = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_chat_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_chat_language_model: str = "gpt-4.1-mini" - openai_extract_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_extract_api_base: str = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_extract_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_extract_language_model: str = "gpt-4.1-mini" - openai_text2gql_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_text2gql_api_base: str = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_text2gql_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_text2gql_language_model: str = "gpt-4.1-mini" - openai_embedding_api_base: str = os.environ.get("OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1") + openai_embedding_api_base: str = os.environ.get( + "OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1" + ) openai_embedding_api_key: str | None = os.environ.get("OPENAI_EMBEDDING_API_KEY") openai_embedding_model: str = "text-embedding-3-small" openai_embedding_model_dim: int = 1536 @@ -49,7 +57,9 @@ class LLMConfig(BaseConfig): openai_extract_tokens: int = 256 openai_text2gql_tokens: int = 4096 # 2. Rerank settings - cohere_base_url: str = os.environ.get("CO_API_URL", "https://api.cohere.com/v1/rerank") + cohere_base_url: str = os.environ.get( + "CO_API_URL", "https://api.cohere.com/v1/rerank" + ) reranker_api_key: Optional[str] = None reranker_model: Optional[str] = None # 3. Ollama settings diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py index 5fec3a778..4ec9256c5 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py @@ -24,7 +24,9 @@ from hugegraph_llm.utils.log import log dir_name = os.path.dirname -env_path = os.path.join(os.getcwd(), ".env") # Load .env from the current working directory +env_path = os.path.join( + os.getcwd(), ".env" +) # Load .env from the current working directory class BaseConfig(BaseSettings): @@ -88,7 +90,9 @@ def check_env(self): # Step 2: Add missing config items to .env self._sync_object_to_env(env_config, config_dict) except Exception as e: - log.error("An error occurred when checking the .env variable file: %s", str(e)) + log.error( + "An error occurred when checking the .env variable file: %s", str(e) + ) raise def _sync_env_to_object(self, env_config, config_dict): @@ -139,7 +143,9 @@ def __init__(self, **data): # Synchronize configurations between the object and .env file self.check_env() - log.info("The %s file was loaded. Class: %s", env_path, self.__class__.__name__) + log.info( + "The %s file was loaded. Class: %s", env_path, self.__class__.__name__ + ) except Exception as e: log.error("An error occurred when initializing the configuration object: %s", str(e)) raise diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py index 2369d01a6..b15bad0a6 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py @@ -57,7 +57,9 @@ def ensure_yaml_file_exists(self): current_dir = Path.cwd().resolve() project_root = get_project_root() if current_dir == project_root: - log.info("Current working directory is the project root, proceeding to run the app.") + log.info( + "Current working directory is the project root, proceeding to run the app." + ) else: error_msg = ( f"Current working directory is not the project root. " @@ -122,7 +124,9 @@ def to_literal(val): "gremlin_generate_prompt": to_literal(self.gremlin_generate_prompt), "doc_input_text": to_literal(self.doc_input_text), "_language_generated": str(self.llm_settings.language).lower().strip(), - "generate_extract_prompt_template": to_literal(self.generate_extract_prompt_template), + "generate_extract_prompt_template": to_literal( + self.generate_extract_prompt_template + ), } with open(yaml_file_path, "w", encoding="utf-8") as file: yaml.dump(data, file, allow_unicode=True, sort_keys=False, default_flow_style=False) @@ -150,7 +154,9 @@ def generate_yaml_file(self): self.keywords_extract_prompt = self.keywords_extract_prompt_EN self.doc_input_text = self.doc_input_text_EN self.save_to_yaml() - log.info("Prompt file '%s' has been generated with default values.", yaml_file_path) + log.info( + "Prompt file '%s' has been generated with default values.", yaml_file_path + ) def update_yaml_file(self): self.save_to_yaml() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 0b041387e..d3beebcbb 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -69,7 +69,7 @@ def clear_llm_server_log(): # Function to validate password and control access to logs -def check_password(password, request = None): +def check_password(password, request=None): client_ip = request.client.host if request else "Unknown IP" admin_token = admin_settings.admin_token @@ -109,6 +109,7 @@ def create_admin_block(): # Error message box, initially hidden <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD error_message = gr.Textbox( label="", visible=False, interactive=False, elem_classes="error-message" @@ -119,6 +120,11 @@ def create_admin_block(): ======= error_message = gr.Textbox(label="", visible=False, interactive=False, elem_classes="error-message") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + error_message = gr.Textbox( + label="", visible=False, interactive=False, elem_classes="error-message" + ) +>>>>>>> 3aeef7d (fix) # Button to submit password submit_button = gr.Button("Submit") @@ -137,7 +143,9 @@ def create_admin_block(): with gr.Row(): with gr.Column(): # Button to clear LLM Server log, initially hidden - clear_llm_server_button = gr.Button("Clear LLM Server Log", visible=False) + clear_llm_server_button = gr.Button( + "Clear LLM Server Log", visible=False + ) with gr.Column(): # Button to refresh LLM Server log manually refresh_llm_server_button = gr.Button( 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 a8450ab5f..a78e62361 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -93,7 +93,9 @@ def init_rag_ui() -> gr.Interface: textbox_array_graph_config = create_configs_block() with gr.Tab(label="1. Build RAG Index 💡"): - textbox_input_text, textbox_input_schema, textbox_info_extract_template = create_vector_graph_block() + textbox_input_text, textbox_input_schema, textbox_info_extract_template = ( + create_vector_graph_block() + ) with gr.Tab(label="2. (Graph)RAG & User Functions 📖"): ( textbox_inp, @@ -102,7 +104,9 @@ def init_rag_ui() -> gr.Interface: textbox_custom_related_information, ) = create_rag_block() with gr.Tab(label="3. Text2gremlin ⚙️"): - textbox_gremlin_inp, textbox_gremlin_schema, textbox_gremlin_prompt = create_text2gremlin_block() + textbox_gremlin_inp, textbox_gremlin_schema, textbox_gremlin_prompt = ( + create_text2gremlin_block() + ) with gr.Tab(label="4. Graph Tools 🚧"): create_other_block() with gr.Tab(label="5. Admin Tools 🛠"): @@ -162,7 +166,9 @@ def create_app(): prompt.update_yaml_file() assert admin_settings.enable_login auth_enabled = admin_settings.enable_login.lower() == "true" - log.info("(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled") + log.info( + "(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled" + ) api_auth = APIRouter(dependencies=[Depends(authenticate)] if auth_enabled else []) hugegraph_llm = init_rag_ui() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 8d90da414..155a4b0b8 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -65,12 +65,16 @@ def test_litellm_chat(api_key, api_base, model_name, max_tokens: int) -> int: return 200 -def test_api_connection(url, method="GET", headers=None, params=None, body=None, auth=None, origin_call=None) -> int: +def test_api_connection( + url, method="GET", headers=None, params=None, body=None, auth=None, origin_call=None +) -> int: # TODO: use fastapi.request / starlette instead? log.debug("Request URL: %s", url) try: if method.upper() == "GET": - resp = requests.get(url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth) + resp = requests.get( + url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth + ) elif method.upper() == "POST": resp = requests.post( url, @@ -105,10 +109,16 @@ def test_api_connection(url, method="GET", headers=None, params=None, body=None, return resp.status_code +<<<<<<< HEAD <<<<<<< HEAD def apply_embedding_config(arg1, arg2, arg3, origin_call=None) -> int: ======= def config_qianfan_model(arg1, arg2, arg3=None, settings_prefix=None, origin_call=None) -> int: +======= +def config_qianfan_model( + arg1, arg2, arg3=None, settings_prefix=None, origin_call=None +) -> int: +>>>>>>> 3aeef7d (fix) setattr(llm_settings, f"qianfan_{settings_prefix}_api_key", arg1) setattr(llm_settings, f"qianfan_{settings_prefix}_secret_key", arg2) if arg3: @@ -138,6 +148,7 @@ def apply_embedding_config(arg1, arg2, arg3, arg4, origin_call=None) -> int: headers = {"Authorization": f"Bearer {arg1}"} data = {"model": arg3, "input": "test"} <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD status_code = test_api_connection( test_url, method="POST", headers=headers, body=data, origin_call=origin_call @@ -152,12 +163,19 @@ def apply_embedding_config(arg1, arg2, arg3, arg4, origin_call=None) -> int: ======= status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + status_code = test_api_connection( + test_url, method="POST", headers=headers, body=data, origin_call=origin_call + ) +>>>>>>> 3aeef7d (fix) elif embedding_option == "ollama/local": llm_settings.ollama_embedding_host = arg1 llm_settings.ollama_embedding_port = int(arg2) llm_settings.ollama_embedding_model = arg3 llm_settings.ollama_embedding_model_dim = arg4 - status_code = test_api_connection(f"http://{arg1}:{arg2}", origin_call=origin_call) + status_code = test_api_connection( + f"http://{arg1}:{arg2}", origin_call=origin_call + ) elif embedding_option == "litellm": llm_settings.litellm_embedding_api_key = arg1 llm_settings.litellm_embedding_api_base = arg2 @@ -248,6 +266,7 @@ def apply_llm_config( setattr(llm_settings, f"openai_{current_llm_config}_language_model", model_name) setattr(llm_settings, f"openai_{current_llm_config}_tokens", int(max_tokens)) +<<<<<<< HEAD <<<<<<< HEAD <<<<<<< HEAD test_url = ( @@ -259,6 +278,12 @@ def apply_llm_config( ======= test_url = getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + test_url = ( + getattr(llm_settings, f"openai_{current_llm_config}_api_base") + + "/chat/completions" + ) +>>>>>>> 3aeef7d (fix) data = { "model": model_name, "temperature": 0.01, @@ -267,6 +292,7 @@ def apply_llm_config( <<<<<<< HEAD headers = {"Authorization": f"Bearer {api_key_or_host}"} <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD status_code = test_api_connection( test_url, method="POST", headers=headers, body=data, origin_call=origin_call @@ -274,6 +300,11 @@ def apply_llm_config( ======= headers = {"Authorization": f"Bearer {arg1}"} status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) +======= + status_code = test_api_connection( + test_url, method="POST", headers=headers, body=data, origin_call=origin_call + ) +>>>>>>> 3aeef7d (fix) elif llm_option == "qianfan_wenxin": status_code = config_qianfan_model( @@ -289,17 +320,27 @@ def apply_llm_config( elif llm_option == "ollama/local": setattr(llm_settings, f"ollama_{current_llm_config}_host", api_key_or_host) - setattr(llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port)) + setattr( + llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port) + ) setattr(llm_settings, f"ollama_{current_llm_config}_language_model", model_name) - status_code = test_api_connection(f"http://{api_key_or_host}:{api_base_or_port}", origin_call=origin_call) + status_code = test_api_connection( + f"http://{api_key_or_host}:{api_base_or_port}", origin_call=origin_call + ) elif llm_option == "litellm": setattr(llm_settings, f"litellm_{current_llm_config}_api_key", api_key_or_host) - setattr(llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port) - setattr(llm_settings, f"litellm_{current_llm_config}_language_model", model_name) + setattr( + llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port + ) + setattr( + llm_settings, f"litellm_{current_llm_config}_language_model", model_name + ) setattr(llm_settings, f"litellm_{current_llm_config}_tokens", int(max_tokens)) - status_code = test_litellm_chat(api_key_or_host, api_base_or_port, model_name, int(max_tokens)) + status_code = test_litellm_chat( + api_key_or_host, api_base_or_port, model_name, int(max_tokens) + ) gr.Info("Configured!") llm_settings.update_env() @@ -368,7 +409,9 @@ def create_configs_block() -> list: ), ] graph_config_button = gr.Button("Apply Configuration") - graph_config_button.click(apply_graph_config, inputs=graph_config_input) # pylint: disable=no-member + graph_config_button.click( + apply_graph_config, inputs=graph_config_input + ) # pylint: disable=no-member # TODO : use OOP to refactor the following code with gr.Accordion("2. Set up the LLM.", open=False): @@ -520,18 +563,27 @@ def chat_llm_settings(llm_type): >>>>>>> f42fa9b (feat(llm): use lambda) ] else: - llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] + llm_config_input = [ + gr.Textbox(value="", visible=False) for _ in range(4) + ] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click(apply_llm_config_with_chat_op, inputs=llm_config_input) + llm_config_button.click( + apply_llm_config_with_chat_op, inputs=llm_config_input + ) # Determine whether there are Settings in the.env file - env_path = os.path.join(os.getcwd(), ".env") # Load .env from the current working directory + env_path = os.path.join( + os.getcwd(), ".env" + ) # Load .env from the current working directory env_vars = dotenv_values(env_path) api_extract_key = env_vars.get("OPENAI_EXTRACT_API_KEY") api_text2sql_key = env_vars.get("OPENAI_TEXT2GQL_API_KEY") if not api_extract_key: - llm_config_button.click(apply_llm_config_with_text2gql_op, inputs=llm_config_input) + llm_config_button.click( + apply_llm_config_with_text2gql_op, inputs=llm_config_input + ) if not api_text2sql_key: <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD llm_config_button.click( apply_llm_config_with_extract_op, inputs=llm_config_input @@ -539,6 +591,11 @@ def chat_llm_settings(llm_type): ======= llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + llm_config_button.click( + apply_llm_config_with_extract_op, inputs=llm_config_input + ) +>>>>>>> 3aeef7d (fix) with gr.Tab(label="mini_tasks"): extract_llm_dropdown = gr.Dropdown( @@ -575,7 +632,9 @@ def extract_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr(llm_settings, "openai_extract_language_model"), + value=getattr( + llm_settings, "openai_extract_language_model" + ), label="model_name", ), gr.Textbox( @@ -594,7 +653,9 @@ def extract_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr(llm_settings, "ollama_extract_language_model"), + value=getattr( + llm_settings, "ollama_extract_language_model" + ), label="model_name", ), ======= @@ -669,7 +730,9 @@ def extract_llm_settings(llm_type): info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr(llm_settings, "litellm_extract_language_model"), + value=getattr( + llm_settings, "litellm_extract_language_model" + ), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -687,9 +750,13 @@ def extract_llm_settings(llm_type): >>>>>>> f42fa9b (feat(llm): use lambda) ] else: - llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] + llm_config_input = [ + gr.Textbox(value="", visible=False) for _ in range(4) + ] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) + llm_config_button.click( + apply_llm_config_with_extract_op, inputs=llm_config_input + ) <<<<<<< HEAD <<<<<<< HEAD @@ -728,7 +795,9 @@ def text2gql_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr(llm_settings, "openai_text2gql_language_model"), + value=getattr( + llm_settings, "openai_text2gql_language_model" + ), label="model_name", ), gr.Textbox( @@ -747,7 +816,9 @@ def text2gql_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr(llm_settings, "ollama_text2gql_language_model"), + value=getattr( + llm_settings, "ollama_text2gql_language_model" + ), label="model_name", ), ======= @@ -799,6 +870,7 @@ def text2gql_llm_settings(llm_type): llm_config_input = [ gr.Textbox( <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "litellm_text2gql_api_key"), label="api_key", @@ -808,24 +880,39 @@ def text2gql_llm_settings(llm_type): >>>>>>> 38dce0b (feat(llm): vector db finished) ======= value=lambda: getattr(llm_settings, "litellm_text2gql_api_key"), +======= + value=lambda: getattr( + llm_settings, "litellm_text2gql_api_key" + ), +>>>>>>> 3aeef7d (fix) label="api_key", type="password", >>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_text2gql_api_base"), + value=lambda: getattr( + llm_settings, "litellm_text2gql_api_base" + ), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_text2gql_language_model"), + value=lambda: getattr( + llm_settings, "litellm_text2gql_language_model" + ), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), <<<<<<< HEAD <<<<<<< HEAD gr.Textbox( +<<<<<<< HEAD value=getattr(llm_settings, "litellm_text2gql_tokens"), +======= + value=lambda: getattr( + llm_settings, "litellm_text2gql_tokens" + ), +>>>>>>> 3aeef7d (fix) label="max_token", ), ======= @@ -836,9 +923,13 @@ def text2gql_llm_settings(llm_type): >>>>>>> f42fa9b (feat(llm): use lambda) ] else: - llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] + llm_config_input = [ + gr.Textbox(value="", visible=False) for _ in range(4) + ] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click(apply_llm_config_with_text2gql_op, inputs=llm_config_input) + llm_config_button.click( + apply_llm_config_with_text2gql_op, inputs=llm_config_input + ) with gr.Accordion("3. Set up the Embedding.", open=False): embedding_dropdown = gr.Dropdown( @@ -931,6 +1022,7 @@ def embedding_settings(embedding_type): embedding_config_input = [ gr.Textbox( <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD value=getattr(llm_settings, "litellm_embedding_api_key"), label="api_key", @@ -940,24 +1032,35 @@ def embedding_settings(embedding_type): >>>>>>> 38dce0b (feat(llm): vector db finished) ======= value=lambda: getattr(llm_settings, "litellm_embedding_api_key"), +======= + value=lambda: getattr( + llm_settings, "litellm_embedding_api_key" + ), +>>>>>>> 3aeef7d (fix) label="api_key", type="password", >>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_embedding_api_base"), + value=lambda: getattr( + llm_settings, "litellm_embedding_api_base" + ), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_embedding_model"), + value=lambda: getattr( + llm_settings, "litellm_embedding_model" + ), label="model_name", info="Please refer to https://docs.litellm.ai/docs/embedding/supported_embedding", ), <<<<<<< HEAD ======= gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_embedding_model_dim"), + value=lambda: getattr( + llm_settings, "litellm_embedding_model_dim" + ), label="model_dim", type="text", ), @@ -988,7 +1091,9 @@ def embedding_settings(embedding_type): @gr.render(inputs=[reranker_dropdown]) def reranker_settings(reranker_type): - llm_settings.reranker_type = reranker_type if reranker_type != "None" else None + llm_settings.reranker_type = ( + reranker_type if reranker_type != "None" else None + ) if reranker_type == "cohere": with gr.Row(): reranker_config_input = [ @@ -998,6 +1103,7 @@ def reranker_settings(reranker_type): label="api_key", type="password", ), +<<<<<<< HEAD gr.Textbox(value=llm_settings.reranker_model, label="model"), gr.Textbox(value=llm_settings.cohere_base_url, label="base_url"), ======= @@ -1005,6 +1111,14 @@ def reranker_settings(reranker_type): gr.Textbox(value=lambda: llm_settings.reranker_model, label="model"), gr.Textbox(value=lambda: llm_settings.cohere_base_url, label="base_url"), >>>>>>> f42fa9b (feat(llm): use lambda) +======= + gr.Textbox( + value=lambda: llm_settings.reranker_model, label="model" + ), + gr.Textbox( + value=lambda: llm_settings.cohere_base_url, label="base_url" + ), +>>>>>>> 3aeef7d (fix) ] elif reranker_type == "siliconflow": with gr.Row(): diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py index 82650b907..3f8089b6d 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py @@ -34,9 +34,13 @@ def create_other_block(): inp = gr.Textbox( value="g.V().limit(10)", label="Gremlin query", show_copy_button=True, lines=8 ) - out = gr.Code(label="Output", language="json", elem_classes="code-container-show") + out = gr.Code( + label="Output", language="json", elem_classes="code-container-show" + ) btn = gr.Button("Run Gremlin query") - btn.click(fn=run_gremlin_query, inputs=[inp], outputs=out) # pylint: disable=no-member + btn.click( + fn=run_gremlin_query, inputs=[inp], outputs=out + ) # pylint: disable=no-member gr.Markdown("---") with gr.Row(): @@ -51,7 +55,9 @@ def create_other_block(): inp = [] out = gr.Textbox(label="Init Graph Demo Result", show_copy_button=True) btn = gr.Button("(BETA) Init HugeGraph test data (🚧)") - btn.click(fn=init_hg_test_data, inputs=inp, outputs=out) # pylint: disable=no-member + btn.click( + fn=init_hg_test_data, inputs=inp, outputs=out + ) # pylint: disable=no-member @asynccontextmanager diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 03720d810..de4e0b82c 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -137,7 +137,11 @@ def rag_answer( vector_dis_threshold=vector_dis_threshold, topk_per_keyword=topk_per_keyword, ) +<<<<<<< HEAD if res.get("switch_to_bleu"): +======= + if context.get("switch_to_bleu"): +>>>>>>> 3aeef7d (fix) gr.Warning( "Online reranker fails, automatically switches to local bleu rerank." ) @@ -262,10 +266,20 @@ async def rag_answer_streaming( query=text, vector_search=vector_search, graph_search=graph_search, +<<<<<<< HEAD +======= + ) + if context.get("switch_to_bleu"): + gr.Warning( + "Online reranker fails, automatically switches to local bleu rerank." + ) + answer_synthesize = AnswerSynthesize( +>>>>>>> 3aeef7d (fix) raw_answer=raw_answer, vector_only_answer=vector_only_answer, graph_only_answer=graph_only_answer, graph_vector_answer=graph_vector_answer, +<<<<<<< HEAD graph_ratio=graph_ratio, rerank_method=rerank_method, near_neighbor_first=near_neighbor_first, @@ -276,6 +290,12 @@ async def rag_answer_streaming( gremlin_prompt=gremlin_prompt, ): if res.get("switch_to_bleu"): +======= + prompt_template=answer_prompt, + ) + async for context in answer_synthesize.run_streaming(context): + if context.get("switch_to_bleu"): +>>>>>>> 3aeef7d (fix) gr.Warning( "Online reranker fails, automatically switches to local bleu rerank." ) @@ -349,19 +369,29 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> 3aeef7d (fix) raw_radio = gr.Radio( choices=[True, False], value=False, label="Basic LLM Answer" ) vector_only_radio = gr.Radio( choices=[True, False], value=False, label="Vector-only Answer" ) +<<<<<<< HEAD ======= raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") vector_only_radio = gr.Radio(choices=[True, False], value=False, label="Vector-only Answer") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= +>>>>>>> 3aeef7d (fix) with gr.Row(): - graph_only_radio = gr.Radio(choices=[True, False], value=True, label="Graph-only Answer") - graph_vector_radio = gr.Radio(choices=[True, False], value=False, label="Graph-Vector Answer") + graph_only_radio = gr.Radio( + choices=[True, False], value=True, label="Graph-only Answer" + ) + graph_vector_radio = gr.Radio( + choices=[True, False], value=False, label="Graph-Vector Answer" + ) def toggle_slider(enable): return gr.update(interactive=enable) @@ -379,9 +409,13 @@ def toggle_slider(enable): label="Template Num (<0 means disable text2gql) ", precision=0, ) - graph_ratio = gr.Slider(0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False) + graph_ratio = gr.Slider( + 0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False + ) - graph_vector_radio.change(toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio) # pylint: disable=no-member + graph_vector_radio.change( + toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio + ) # pylint: disable=no-member near_neighbor_first = gr.Checkbox( value=False, label="Near neighbor first(Optional)", @@ -510,7 +544,9 @@ def several_rag_answer( with gr.Row(): with gr.Column(): - questions_file = gr.File(file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)") + questions_file = gr.File( + file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" + ) with gr.Column(): test_template_file = os.path.join( resource_path, "demo", "questions_template.xlsx" diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index f3119f67c..1d4407519 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -123,7 +123,16 @@ def build_example_vector_index(temp_file) -> dict: timestamp = datetime.now().strftime("%Y%m%d%H%M%S") _, file_name = os.path.split(f"{name}_{timestamp}{ext}") log.info("Copying file to: %s", file_name) +<<<<<<< HEAD target_file = os.path.join(resource_path, folder_name, "gremlin_examples", file_name) +======= + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + target_file = os.path.join( + resource_path, folder_name, "gremlin_examples", file_name + ) +>>>>>>> 3aeef7d (fix) try: import shutil @@ -232,16 +241,26 @@ def gremlin_generate( vector_index = get_vector_index_class(index_settings.now_vector_index) ======= vector_index = get_vector_index_class(index_settings.cur_vector_index) +<<<<<<< HEAD >>>>>>> a255aed (fix cycle import & add docs) generator = GremlinGenerator(llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding()) >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + generator = GremlinGenerator( + llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() + ) +>>>>>>> 3aeef7d (fix) sm = SchemaManager(graph_name=schema) processed_schema, short_schema = _process_schema(schema, generator, sm) if processed_schema is None and short_schema is None: - return GremlinResult.error("Invalid JSON schema, please check the format carefully.") + return GremlinResult.error( + "Invalid JSON schema, please check the format carefully." + ) - updated_schema = sm.simple_schema(processed_schema) if short_schema else processed_schema + updated_schema = ( + sm.simple_schema(processed_schema) if short_schema else processed_schema + ) store_schema(str(updated_schema), inp, gremlin_prompt) output_types = _configure_output_types(requested_outputs) @@ -254,7 +273,9 @@ def gremlin_generate( _execute_queries(context, output_types) - match_result = json.dumps(context.get("match_result", "No Results"), ensure_ascii=False, indent=2) + match_result = json.dumps( + context.get("match_result", "No Results"), ensure_ascii=False, indent=2 + ) return GremlinResult.success_result( match_result=match_result, template_gremlin=context["result"], @@ -271,14 +292,22 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} + new_vertex = { + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex + } mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) if "edgelabels" in schema: mini_schema["edgelabels"] = [] for edge in schema["edgelabels"]: - new_edge = {key: edge[key] for key in ["name", "source_label", "target_label", "properties"] if key in edge} + new_edge = { + key: edge[key] + for key in ["name", "source_label", "target_label", "properties"] + if key in edge + } mini_schema["edgelabels"].append(new_edge) return mini_schema @@ -349,7 +378,9 @@ def create_text2gremlin_block() -> Tuple: with gr.Row(): btn = gr.Button("Build Example Vector Index", variant="primary") - btn.click(build_example_vector_index, inputs=[file], outputs=[out]) # pylint: disable=no-member + btn.click( + build_example_vector_index, inputs=[file], outputs=[out] + ) # pylint: disable=no-member gr.Markdown("## Nature Language To Gremlin") with gr.Row(): @@ -362,8 +393,12 @@ def create_text2gremlin_block() -> Tuple: language="javascript", elem_classes="code-container-show", ) - initialized_out = gr.Textbox(label="Gremlin With Template", show_copy_button=True) - raw_out = gr.Textbox(label="Gremlin Without Template", show_copy_button=True) + initialized_out = gr.Textbox( + label="Gremlin With Template", show_copy_button=True + ) + raw_out = gr.Textbox( + label="Gremlin Without Template", show_copy_button=True + ) tmpl_exec_out = gr.Code( label="Query With Template Output", language="json", @@ -376,7 +411,9 @@ def create_text2gremlin_block() -> Tuple: ) with gr.Column(scale=1): - example_num_slider = gr.Slider(minimum=0, maximum=10, step=1, value=2, label="Number of refer examples") + example_num_slider = gr.Slider( + minimum=0, maximum=10, step=1, value=2, label="Number of refer examples" + ) schema_box = gr.Textbox( value=prompt.text2gql_graph_schema, label="Schema", lines=2, show_copy_button=True ) @@ -464,7 +501,9 @@ def gremlin_generate_selective( if not requested_outputs: # None or empty list requested_outputs = output_keys - result = gremlin_generate(inp, example_num, schema_input, gremlin_prompt_input, requested_outputs) + result = gremlin_generate( + inp, example_num, schema_input, gremlin_prompt_input, requested_outputs + ) outputs_dict: Dict[str, Any] = {} diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index af6570b19..92ab0b895 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -58,7 +58,11 @@ def store_prompt(doc, schema, example_prompt): # update env variables: doc, schema and example_prompt - if prompt.doc_input_text != doc or prompt.graph_schema != schema or prompt.extract_graph_prompt != example_prompt: + if ( + prompt.doc_input_text != doc + or prompt.graph_schema != schema + or prompt.extract_graph_prompt != example_prompt + ): prompt.doc_input_text = doc prompt.graph_schema = schema prompt.extract_graph_prompt = example_prompt @@ -86,7 +90,9 @@ def generate_prompt_for_ui(source_text, scenario, example_name): def load_example_names(): """Load all candidate examples""" try: - examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "prompt_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return [example.get("name", "Unnamed example") for example in examples] @@ -102,29 +108,41 @@ def load_query_examples(): "language", <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> 3aeef7d (fix) ( getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") else "EN" ), +<<<<<<< HEAD ======= getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") else "EN", >>>>>>> 87ee5d3 (style: format code with black line-length 120) ======= (getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") else "EN"), >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= +>>>>>>> 3aeef7d (fix) ) if language.upper() == "CN": - examples_path = os.path.join(resource_path, "prompt_examples", "query_examples_CN.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "query_examples_CN.json" + ) else: - examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "query_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) except (FileNotFoundError, json.JSONDecodeError): try: - examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "query_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -135,7 +153,9 @@ def load_query_examples(): def load_schema_fewshot_examples(): """Load few-shot examples from a JSON file""" try: - examples_path = os.path.join(resource_path, "prompt_examples", "schema_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "schema_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -557,7 +577,9 @@ def create_vector_graph_block(): max_lines=29, ) - out = gr.Code(label="Output Info", language="json", elem_classes="code-container-edit") + out = gr.Code( + label="Output Info", language="json", elem_classes="code-container-edit" + ) with gr.Row(): with gr.Accordion("Get RAG Info", open=False): @@ -606,7 +628,9 @@ def create_vector_graph_block(): store_prompt, inputs=[input_text, input_schema, info_extract_template], ) - vector_import_bt.click(build_vector_index, inputs=[input_file, input_text], outputs=out).then( + vector_import_bt.click( + build_vector_index, inputs=[input_file, input_text], outputs=out + ).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) @@ -634,15 +658,17 @@ def create_vector_graph_block(): inputs=[input_text, input_schema, info_extract_template], ) - graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( - update_vid_embedding - ).then( + graph_loading_bt.click( + import_graph_data, inputs=[out, input_schema], outputs=[out] + ).then(update_vid_embedding).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) build_schema_bt.click( - lambda it, qe, fs: extract_graph([], it, prompt.graph_schema, prompt.extract_graph_prompt), + lambda it, qe, fs: extract_graph( + [], it, prompt.graph_schema, prompt.extract_graph_prompt + ), inputs=[input_text, query_example, few_shot], outputs=[input_schema], ).then( diff --git a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py index 5012f7e01..11369c06a 100644 --- a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py +++ b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py @@ -38,6 +38,7 @@ def __init__( raise ValueError("Argument `language` must be zh or en!") if split_type == "paragraph": <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=500, chunk_overlap=30, separators=separators @@ -56,6 +57,15 @@ def __init__( elif split_type == "sentence": self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=50, chunk_overlap=0, separators=separators) >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + self.text_splitter = RecursiveCharacterTextSplitter( + chunk_size=500, chunk_overlap=30, separators=separators + ) + elif split_type == "sentence": + self.text_splitter = RecursiveCharacterTextSplitter( + chunk_size=50, chunk_overlap=0, separators=separators + ) +>>>>>>> 3aeef7d (fix) else: raise ValueError("Arg `type` must be paragraph, sentence!") diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py index feda7a24c..2e1cfc267 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py @@ -56,7 +56,9 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: """ @abstractmethod - def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9) -> List[Any]: + def search( + self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9 + ) -> List[Any]: """ Search for the top_k most similar vectors to the query vector. diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py index d0a016c53..a8f23155a 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py @@ -66,7 +66,9 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: self.properties = [p for i, p in enumerate(self.properties) if i not in indices] return remove_num - def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9) -> List[Any]: + def search( + self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9 + ) -> List[Any]: if self.index.ntotal == 0: return [] diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py index 631e7908b..fae5ca860 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -73,9 +73,15 @@ def __init__( def _create_collection(self): """Create a new collection in Milvus.""" - id_field = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True) - vector_field = FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embed_dim) - property_field = FieldSchema(name="property", dtype=DataType.VARCHAR, max_length=65535) + id_field = FieldSchema( + name="id", dtype=DataType.INT64, is_primary=True, auto_id=True + ) + vector_field = FieldSchema( + name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embed_dim + ) + property_field = FieldSchema( + name="property", dtype=DataType.VARCHAR, max_length=65535 + ) original_id_field = FieldSchema(name="original_id", dtype=DataType.INT64) schema = CollectionSchema( @@ -147,7 +153,9 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: finally: self.collection.release() - def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9) -> List[Any]: + def search( + self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9 + ) -> List[Any]: try: if self.collection.num_entities == 0: return [] diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index 7e403f312..14b97fbff 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -28,7 +28,9 @@ class QdrantVectorIndex(VectorStoreBase): - def __init__(self, name: str, host: str, port: int, api_key=None, embed_dim: int = 1024): + def __init__( + self, name: str, host: str, port: int, api_key=None, embed_dim: int = 1024 + ): self.embed_dim = embed_dim self.host = host self.port = port @@ -55,7 +57,9 @@ def _create_collection(self): """Create a new collection in Qdrant.""" self.client.create_collection( collection_name=self.name, - vectors_config=models.VectorParams(size=self.embed_dim, distance=models.Distance.COSINE), + vectors_config=models.VectorParams( + size=self.embed_dim, distance=models.Distance.COSINE + ), ) log.info("Created Qdrant collection '%s'", self.name) @@ -113,8 +117,12 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: return remove_num - def search(self, query_vector: List[float], top_k: int = 5, dis_threshold: float = 0.9): - search_result = self.client.search(collection_name=self.name, query_vector=query_vector, limit=top_k) + def search( + self, query_vector: List[float], top_k: int = 5, dis_threshold: float = 0.9 + ): + search_result = self.client.search( + collection_name=self.name, query_vector=query_vector, limit=top_k + ) result_properties = [] diff --git a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py index f13e11e6f..5d98ebdf8 100644 --- a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py +++ b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py @@ -41,7 +41,9 @@ async def dispatch(self, request: Request, call_next): unit = "s" response.headers["X-Process-Time"] = f"{process_time:.2f} {unit}" - log.info("Request process time: %.2f ms, code=%d", process_time, response.status_code) + log.info( + "Request process time: %.2f ms, code=%d", process_time, response.status_code + ) log.info( "%s - Args: %s, IP: %s, URL: %s", request.method, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index d96840911..8e6af2774 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -65,7 +65,9 @@ def __init__(self): def get_embedding(self): if self.embedding_type == "openai": - assert llm_settings.openai_embedding_model_dim, "openai_embedding_model_dim is need" + assert ( + llm_settings.openai_embedding_model_dim + ), "openai_embedding_model_dim is need" return OpenAIEmbedding( embedding_dimension=llm_settings.openai_embedding_model_dim, model_name=llm_settings.openai_embedding_model, @@ -73,7 +75,9 @@ def get_embedding(self): api_base=llm_settings.openai_embedding_api_base, ) if self.embedding_type == "ollama/local": - assert llm_settings.ollama_embedding_model_dim, "ollama_embedding_model_dim is need" + assert ( + llm_settings.ollama_embedding_model_dim + ), "ollama_embedding_model_dim is need" return OllamaEmbedding( <<<<<<< HEAD model_name=llm_settings.ollama_embedding_model, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index 6e30cca71..15d928286 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -91,5 +91,7 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] A list of embedding vectors, where each vector is a list of floats. The order of embeddings should match the order of input texts. """ - response = await self.aclient.embeddings.create(input=texts, model=self.model_name) + response = await self.aclient.embeddings.create( + input=texts, model=self.model_name + ) return [data.embedding for data in response.data] 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 9121fca09..7e1eaab68 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py @@ -173,4 +173,8 @@ def get_text2gql_llm(self): if __name__ == "__main__": client = LLMs().get_chat_llm() print(client.generate(prompt="What is the capital of China?")) - print(client.generate(messages=[{"role": "user", "content": "What is the capital of China?"}])) + print( + client.generate( + messages=[{"role": "user", "content": "What is the capital of China?"}] + ) + ) diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py index 2fa979ed9..c15c5440e 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py @@ -29,7 +29,9 @@ class OllamaClient(BaseLLM): """LLM wrapper should take in a prompt and return a string.""" - def __init__(self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwargs): + def __init__( + self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwargs + ): self.model = model self.client = ollama.Client(host=f"http://{host}:{port}", **kwargs) self.async_client = ollama.AsyncClient(host=f"http://{host}:{port}", **kwargs) @@ -99,7 +101,9 @@ def generate_streaming( for chunk in self.client.chat(model=self.model, messages=messages, stream=True): if not chunk["message"]: - log.debug("Received empty chunk['message'] in streaming chunk: %s", chunk) + log.debug( + "Received empty chunk['message'] in streaming chunk: %s", chunk + ) continue token = chunk["message"]["content"] if on_token_callback: @@ -119,6 +123,7 @@ async def agenerate_streaming( try: <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD async_generator = await self.async_client.chat( model=self.model, messages=messages, stream=True @@ -129,6 +134,11 @@ async def agenerate_streaming( ======= async_generator = await self.async_client.chat(model=self.model, messages=messages, stream=True) >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + async_generator = await self.async_client.chat( + model=self.model, messages=messages, stream=True + ) +>>>>>>> 3aeef7d (fix) async for chunk in async_generator: token = chunk.get("message", {}).get("content", "") if on_token_callback: diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py index e1088c890..52d624941 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py @@ -52,7 +52,9 @@ def __init__( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), + retry=retry_if_exception_type( + (RateLimitError, APIConnectionError, APITimeoutError) + ), ) def generate( self, @@ -87,7 +89,9 @@ def generate( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), + retry=retry_if_exception_type( + (RateLimitError, APIConnectionError, APITimeoutError) + ), ) async def agenerate( self, @@ -122,7 +126,9 @@ async def agenerate( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), + retry=retry_if_exception_type( + (RateLimitError, APIConnectionError, APITimeoutError) + ), ) def generate_streaming( self, diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py index fd4643e3c..9886aa0ca 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py @@ -31,10 +31,14 @@ def __init__( self.base_url = base_url self.model = model - def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int] = None) -> List[str]: + def get_rerank_lists( + self, query: str, documents: List[str], top_n: Optional[int] = None + ) -> List[str]: if not top_n: top_n = len(documents) - assert top_n <= len(documents), "'top_n' should be less than or equal to the number of documents" + assert top_n <= len( + documents + ), "'top_n' should be less than or equal to the number of documents" if top_n == 0: return [] @@ -53,7 +57,9 @@ def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int "top_n": top_n, "documents": documents, } - response = requests.post(url, headers=headers, json=payload, timeout=(1.0, 10.0)) + response = requests.post( + url, headers=headers, json=payload, timeout=(1.0, 10.0) + ) response.raise_for_status() # Raise an error for bad status codes results = response.json()["results"] sorted_docs = [documents[item["index"]] for item in results] diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py index aa9f0c061..6136d61b4 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py @@ -32,5 +32,7 @@ def get_reranker(self): model=llm_settings.reranker_model, ) if self.reranker_type == "siliconflow": - return SiliconReranker(api_key=llm_settings.reranker_api_key, model=llm_settings.reranker_model) + return SiliconReranker( + api_key=llm_settings.reranker_api_key, model=llm_settings.reranker_model + ) raise Exception("Reranker type is not supported!") diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py index a67a6ef25..da8a9f7b7 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py @@ -29,10 +29,14 @@ def __init__( self.api_key = api_key self.model = model - def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int] = None) -> List[str]: + def get_rerank_lists( + self, query: str, documents: List[str], top_n: Optional[int] = None + ) -> List[str]: if not top_n: top_n = len(documents) - assert top_n <= len(documents), "'top_n' should be less than or equal to the number of documents" + assert top_n <= len( + documents + ), "'top_n' should be less than or equal to the number of documents" if top_n == 0: return [] @@ -54,7 +58,9 @@ def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int "content-type": Constants.HEADER_CONTENT_TYPE, "authorization": f"Bearer {self.api_key}", } - response = requests.post(url, json=payload, headers=headers, timeout=(1.0, 10.0)) + response = requests.post( + url, json=payload, headers=headers, timeout=(1.0, 10.0) + ) response.raise_for_status() # Raise an error for bad status codes results = response.json()["results"] sorted_docs = [documents[item["index"]] for item in results] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index 63618aaec..bd2479817 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -59,8 +59,12 @@ def _validate_schema(self, schema: Dict[str, Any]) -> None: check_type(schema, dict, "Input data is not a dictionary.") if "vertexlabels" not in schema or "edgelabels" not in schema: log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") - check_type(schema["vertexlabels"], list, "'vertexlabels' in input data is not a list.") - check_type(schema["edgelabels"], list, "'edgelabels' in input data is not a list.") + check_type( + schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." + ) + check_type( + schema["edgelabels"], list, "'edgelabels' in input data is not a list." + ) def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): property_labels = schema.get("propertykeys", []) @@ -72,29 +76,43 @@ def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): property_label_set = {label["name"] for label in property_labels} return property_labels, property_label_set - def _process_vertex_labels(self, schema: Dict[str, Any], property_labels: list, property_label_set: set) -> None: + def _process_vertex_labels( + self, schema: Dict[str, Any], property_labels: list, property_label_set: set + ) -> None: for vertex_label in schema["vertexlabels"]: self._validate_vertex_label(vertex_label) properties = vertex_label["properties"] - primary_keys = self._process_keys(vertex_label, "primary_keys", properties[:1]) + primary_keys = self._process_keys( + vertex_label, "primary_keys", properties[:1] + ) if len(primary_keys) == 0: log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") vertex_label["primary_keys"] = primary_keys - nullable_keys = self._process_keys(vertex_label, "nullable_keys", properties[1:]) + nullable_keys = self._process_keys( + vertex_label, "nullable_keys", properties[1:] + ) vertex_label["nullable_keys"] = nullable_keys - self._add_missing_properties(properties, property_labels, property_label_set) + self._add_missing_properties( + properties, property_labels, property_label_set + ) - def _process_edge_labels(self, schema: Dict[str, Any], property_labels: list, property_label_set: set) -> None: + def _process_edge_labels( + self, schema: Dict[str, Any], property_labels: list, property_label_set: set + ) -> None: for edge_label in schema["edgelabels"]: self._validate_edge_label(edge_label) properties = edge_label.get("properties", []) - self._add_missing_properties(properties, property_labels, property_label_set) + self._add_missing_properties( + properties, property_labels, property_label_set + ) def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") if "name" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'name'.") - check_type(vertex_label["name"], str, "'name' in vertex_label is not of correct type.") + check_type( + vertex_label["name"], str, "'name' in vertex_label is not of correct type." + ) if "properties" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'properties'.") check_type( @@ -107,9 +125,17 @@ def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: check_type(edge_label, dict, "EdgeLabel in input data is not a dictionary.") - if "name" not in edge_label or "source_label" not in edge_label or "target_label" not in edge_label: - log_and_raise("EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'.") - check_type(edge_label["name"], str, "'name' in edge_label is not of correct type.") + if ( + "name" not in edge_label + or "source_label" not in edge_label + or "target_label" not in edge_label + ): + log_and_raise( + "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." + ) + check_type( + edge_label["name"], str, "'name' in edge_label is not of correct type." + ) check_type( edge_label["source_label"], str, @@ -121,13 +147,19 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: "'target_label' in edge_label is not of correct type.", ) - def _process_keys(self, label: Dict[str, Any], key_type: str, default_keys: list) -> list: + def _process_keys( + self, label: Dict[str, Any], key_type: str, default_keys: list + ) -> list: keys = label.get(key_type, default_keys) - check_type(keys, list, f"'{key_type}' in {label['name']} is not of correct type.") + check_type( + keys, list, f"'{key_type}' in {label['name']} is not of correct type." + ) new_keys = [key for key in keys if key in label["properties"]] return new_keys - def _add_missing_properties(self, properties: list, property_labels: list, property_label_set: set) -> None: + def _add_missing_properties( + self, properties: list, property_labels: list, property_label_set: set + ) -> None: for prop in properties: if prop not in property_label_set: property_labels.append( diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py index 743cf3352..a257ccc4c 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py @@ -126,7 +126,8 @@ def _rerank_with_vertex_degree( reranker = Rerankers().get_reranker() try: vertex_rerank_res = [ - reranker.get_rerank_lists(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list + reranker.get_rerank_lists(query, vertex_degree) + [""] + for vertex_degree in vertex_degree_list ] except requests.exceptions.RequestException as e: log.warning( @@ -136,17 +137,25 @@ def _rerank_with_vertex_degree( self.switch_to_bleu = True if self.method == "bleu": - vertex_rerank_res = [_bleu_rerank(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list] + vertex_rerank_res = [ + _bleu_rerank(query, vertex_degree) + [""] + for vertex_degree in vertex_degree_list + ] depth = len(vertex_degree_list) for result in results: if result not in knowledge_with_degree: knowledge_with_degree[result] = [result] + [""] * (depth - 1) if len(knowledge_with_degree[result]) < depth: - knowledge_with_degree[result] += [""] * (depth - len(knowledge_with_degree[result])) + knowledge_with_degree[result] += [""] * ( + depth - len(knowledge_with_degree[result]) + ) def sort_key(res: str) -> Tuple[int, ...]: - return tuple(vertex_rerank_res[i].index(knowledge_with_degree[res][i]) for i in range(depth)) + return tuple( + vertex_rerank_res[i].index(knowledge_with_degree[res][i]) + for i in range(depth) + ) sorted_results = sorted(results, key=sort_key) return sorted_results[:topn] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py index 797ea70ae..c23bf7735 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py @@ -71,7 +71,9 @@ def get_cache_dir() -> str: # Windows (hopefully) else: - local = os.environ.get("LOCALAPPDATA", None) or os.path.expanduser("~\\AppData\\Local") + local = os.environ.get("LOCALAPPDATA", None) or os.path.expanduser( + "~\\AppData\\Local" + ) path = Path(local, "hugegraph_llm") if not os.path.exists(path): 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 a8530632b..c31e77af7 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 @@ -56,7 +56,9 @@ def _get_text_splitter(self, split_type: str): chunk_size=500, chunk_overlap=30, separators=self.separators ).split_text if split_type == SPLIT_TYPE_SENTENCE: - return RecursiveCharacterTextSplitter(chunk_size=50, chunk_overlap=0, separators=self.separators).split_text + return RecursiveCharacterTextSplitter( + chunk_size=50, chunk_overlap=0, separators=self.separators + ).split_text raise ValueError("Type must be paragraph, sentence, html or markdown") def run(self, context: Optional[Dict[str, Any]]) -> Dict[str, Any]: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py index a005c7472..0d160e1e5 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py @@ -77,6 +77,12 @@ def _filter_keywords( results.add(token) sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: - results.update({w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language)}) + results.update( + { + w + for w in sub_tokens + if w not in NLTKHelper().stopwords(lang=self._language) + } + ) return list(results) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index 95eab701f..fa6b79f91 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -42,7 +42,9 @@ class RAGPipeline: querying graph databases and vector indices, merging and re-ranking results, and generating answers. """ - def __init__(self, llm: Optional[BaseLLM] = None, embedding: Optional[BaseEmbedding] = None): + def __init__( + self, llm: Optional[BaseLLM] = None, embedding: Optional[BaseEmbedding] = None + ): """ Initialize the RAGPipeline with optional LLM and embedding models. @@ -252,6 +254,7 @@ def run(self, **kwargs) -> Dict[str, Any]: if len(self._operators) == 0: <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD self.extract_keywords().query_graphdb( max_graph_items=kwargs.get("max_graph_items") @@ -265,6 +268,11 @@ def run(self, **kwargs) -> Dict[str, Any]: ======= self.extract_keywords().query_graphdb(max_graph_items=kwargs.get("max_graph_items")).synthesize_answer() >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + self.extract_keywords().query_graphdb( + max_graph_items=kwargs.get("max_graph_items") + ).synthesize_answer() +>>>>>>> 3aeef7d (fix) context = kwargs diff --git a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py index 7f8c28fa8..52f50fdd6 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py @@ -40,10 +40,16 @@ def clear(self): return self def example_index_build(self, examples, vector_index: type[VectorStoreBase]): - self.operators.append(BuildGremlinExampleIndex(self.embedding, examples, vector_index=vector_index)) + self.operators.append( + BuildGremlinExampleIndex( + self.embedding, examples, vector_index=vector_index + ) + ) return self - def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): + def import_schema( + self, from_hugegraph=None, from_extraction=None, from_user_defined=None + ): if from_hugegraph: self.operators.append(SchemaManager(from_hugegraph)) elif from_user_defined: @@ -55,13 +61,17 @@ def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_def return self def example_index_query(self, num_examples, vector_index: type[VectorStoreBase]): - self.operators.append(GremlinExampleIndexQuery(vector_index, self.embedding, num_examples)) + self.operators.append( + GremlinExampleIndexQuery(vector_index, self.embedding, num_examples) + ) return self def gremlin_generate_synthesize( self, schema, gremlin_prompt: Optional[str] = None, vertices: Optional[List[str]] = None ): - self.operators.append(GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt)) + self.operators.append( + GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt) + ) return self def print_result(self): 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 7aefde57d..aa886c0af 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 @@ -42,13 +42,17 @@ def run(self, data: dict) -> Dict[str, Any]: edges = data.get("edges", []) print(f"get schema {schema}") if not vertices and not edges: - log.critical("(Loading) Both vertices and edges are empty. Please check the input data again.") + log.critical( + "(Loading) Both vertices and edges are empty. Please check the input data again." + ) raise ValueError("Both vertices and edges input are empty.") if not schema: # TODO: ensure the function works correctly (update the logic later) self.schema_free_mode(data.get("triples", [])) - log.warning("Using schema_free mode, could try schema_define mode for better effect!") + log.warning( + "Using schema_free mode, could try schema_define mode for better effect!" + ) else: self.init_schema_if_need(schema) self.load_into_graph(vertices, edges, schema) @@ -64,7 +68,9 @@ def _set_default_property(self, key, input_properties, property_label_map): # list or set default_value = [] input_properties[key] = default_value - log.warning("Property '%s' missing in vertex, set to '%s' for now", key, default_value) + log.warning( + "Property '%s' missing in vertex, set to '%s' for now", key, default_value + ) def _handle_graph_creation(self, func, *args, **kwargs): try: @@ -76,11 +82,17 @@ def _handle_graph_creation(self, func, *args, **kwargs): log.error("Error on creating: %s, %s", args, e) return None - def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many-statements + def load_into_graph( + self, vertices, edges, schema + ): # pylint: disable=too-many-statements # pylint: disable=R0912 (too-many-branches) - vertex_label_map = {v_label["name"]: v_label for v_label in schema["vertexlabels"]} + 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"]} + property_label_map = { + p_label["name"]: p_label for p_label in schema["propertykeys"] + } for vertex in vertices: input_label = vertex["label"] @@ -96,7 +108,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- vertex_label = vertex_label_map[input_label] primary_keys = vertex_label["primary_keys"] nullable_keys = vertex_label.get("nullable_keys", []) - non_null_keys = [key for key in vertex_label["properties"] if key not in nullable_keys] + non_null_keys = [ + key for key in vertex_label["properties"] if key not in nullable_keys + ] has_problem = False # 2. Handle primary-keys mode vertex @@ -128,7 +142,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- # 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) + 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(): @@ -146,7 +162,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- continue # TODO: we could try batch add vertices first, setback to single-mode if failed - vid = self._handle_graph_creation(self.client.graph().addVertex, input_label, input_properties).id + vid = self._handle_graph_creation( + self.client.graph().addVertex, input_label, input_properties + ).id vertex["id"] = vid for edge in edges: @@ -163,7 +181,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- 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) + self._handle_graph_creation( + self.client.graph().addEdge, label, start, end, properties + ) def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -187,19 +207,27 @@ def init_schema_if_need(self, schema: dict): source_vertex_label = edge["source_label"] target_vertex_label = edge["target_label"] properties = edge["properties"] - self.schema.edgeLabel(edge_label).sourceLabel(source_vertex_label).targetLabel( - target_vertex_label - ).properties(*properties).nullableKeys(*properties).ifNotExist().create() + self.schema.edgeLabel(edge_label).sourceLabel( + source_vertex_label + ).targetLabel(target_vertex_label).properties(*properties).nullableKeys( + *properties + ).ifNotExist().create() def schema_free_mode(self, 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( + self.schema.vertexLabel("vertex").useCustomizeStringId().properties( "name" ).ifNotExist().create() + self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel( + "vertex" + ).properties("name").ifNotExist().create() - self.schema.indexLabel("vertexByName").onV("vertex").by("name").secondary().ifNotExist().create() - self.schema.indexLabel("edgeByName").onE("edge").by("name").secondary().ifNotExist().create() + 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: s, p, o = (element.strip() for element in item) @@ -252,7 +280,9 @@ def _set_property_data_type(self, property_key, data_type): log.warning("UUID type is not supported, use text instead") property_key.asText() else: - log.error("Unknown data type %s for property_key %s", data_type, property_key) + log.error( + "Unknown data type %s for property_key %s", data_type, property_key + ) def _set_property_cardinality(self, property_key, cardinality): if cardinality == PropertyCardinality.SINGLE: @@ -262,9 +292,13 @@ def _set_property_cardinality(self, property_key, cardinality): elif cardinality == PropertyCardinality.SET: property_key.valueSet() else: - log.error("Unknown cardinality %s for property_key %s", cardinality, property_key) + 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: + def _check_property_data_type( + self, data_type: str, cardinality: str, value + ) -> bool: if cardinality in ( PropertyCardinality.LIST.value, PropertyCardinality.SET.value, @@ -294,7 +328,9 @@ def _check_single_data_type(self, data_type: str, value) -> bool: if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below - if data_type == PropertyDataType.DATE.value: # the format should be "yyyy-MM-dd" + if ( + data_type == PropertyDataType.DATE.value + ): # the format should be "yyyy-MM-dd" import re return isinstance(value, str) and re.match(r"^\d{4}-\d{2}-\d{2}$", value) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py index 4c4c167c4..73c9530df 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py @@ -47,5 +47,7 @@ def res = [:]; result = self.graph.gremlin().exec(groovy_code)["data"] if isinstance(result, list) and len(result) > 0: - graph_summary.update({key: result[i].get(key) for i, key in enumerate(keys)}) + graph_summary.update( + {key: result[i].get(key) for i, key in enumerate(keys)} + ) return graph_summary diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py index 877a989bf..52399d99e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py @@ -132,7 +132,9 @@ def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: query_embedding = context.get("query_embedding") self._gremlin_generator.clear() - self._gremlin_generator.example_index_query(num_examples=self._num_gremlin_generate_example) + self._gremlin_generator.example_index_query( + num_examples=self._num_gremlin_generate_example + ) gremlin_response = self._gremlin_generator.gremlin_generate_synthesize( context["simple_schema"], vertices=vertices, gremlin_prompt=self._gremlin_prompt ).run(query=query, query_embedding=query_embedding) @@ -146,7 +148,9 @@ def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: result = self._client.gremlin().exec(gremlin=gremlin)["data"] if result == [None]: result = [] - context["graph_result"] = [json.dumps(item, ensure_ascii=False) for item in result] + context["graph_result"] = [ + json.dumps(item, ensure_ascii=False) for item in result + ] if context["graph_result"]: context["graph_result_flag"] = 1 context["graph_context_head"] = ( @@ -196,8 +200,8 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: log.debug("Kneighbor gremlin query: %s", gremlin_query) paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = self._format_graph_query_result( - query_paths=paths + graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( + self._format_graph_query_result(query_paths=paths) ) # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) @@ -220,17 +224,23 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: max_deep=self._max_deep, max_items=self._max_items, ) - log.warning("Unable to find vid, downgraded to property query, please confirm if it meets expectation.") + log.warning( + "Unable to find vid, downgraded to property query, please confirm if it meets expectation." + ) - paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = self._format_graph_query_result( - query_paths=paths + paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)[ + "data" + ] + graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( + self._format_graph_query_result(query_paths=paths) ) context["graph_result"] = list(graph_chain_knowledge) if context["graph_result"]: context["graph_result_flag"] = 0 - context["vertex_degree_list"] = [list(vertex_degree) for vertex_degree in vertex_degree_list] + context["vertex_degree_list"] = [ + list(vertex_degree) for vertex_degree in vertex_degree_list + ] context["knowledge_with_degree"] = knowledge_with_degree context["graph_context_head"] = ( f"The following are graph knowledge in {self._max_deep} depth, e.g:\n" @@ -272,7 +282,9 @@ def _format_graph_from_vertex(self, query_result: List[Any]) -> Set[str]: knowledge.add(node_str) return knowledge - def _format_graph_query_result(self, query_paths) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: + def _format_graph_query_result( + self, query_paths + ) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: use_id_to_match = self._prop_to_match is None subgraph = set() subgraph_with_degree = {} @@ -282,7 +294,9 @@ def _format_graph_query_result(self, query_paths) -> Tuple[Set[str], List[Set[st for path in query_paths: # 1. Process each path - path_str, vertex_with_degree = self._process_path(path, use_id_to_match, v_cache, e_cache) + path_str, vertex_with_degree = self._process_path( + path, use_id_to_match, v_cache, e_cache + ) subgraph.add(path_str) subgraph_with_degree[path_str] = vertex_with_degree # 2. Update vertex degree list @@ -338,13 +352,19 @@ def _process_vertex( use_id_to_match: bool, v_cache: Set[str], ) -> Tuple[str, int, int]: - matched_str = item["id"] if use_id_to_match else item["props"][self._prop_to_match] + matched_str = ( + item["id"] if use_id_to_match else item["props"][self._prop_to_match] + ) if matched_str in node_cache: flat_rel = flat_rel[:-prior_edge_str_len] return flat_rel, prior_edge_str_len, depth node_cache.add(matched_str) - props_str = ", ".join(f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v) + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'v')}" + for k, v in item["props"].items() + if v + ) # TODO: we may remove label id or replace with label name if matched_str in v_cache: @@ -367,10 +387,16 @@ def _process_edge( use_id_to_match: bool, e_cache: Set[Tuple[str, str, str]], ) -> Tuple[str, int]: - props_str = ", ".join(f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v) + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'e')}" + for k, v in item["props"].items() + if v + ) props_str = f"{{{props_str}}}" if props_str else "" prev_matched_str = ( - raw_flat_rel[i - 1]["id"] if use_id_to_match else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] + raw_flat_rel[i - 1]["id"] + if use_id_to_match + else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] ) edge_key = (item["inV"], item["label"], item["outV"]) @@ -380,12 +406,18 @@ def _process_edge( else: edge_label = item["label"] - edge_str = f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" + edge_str = ( + f"--[{edge_label}]-->" + if item["outV"] == prev_matched_str + else f"<--[{edge_label}]--" + ) path_str += edge_str prior_edge_str_len = len(edge_str) return path_str, prior_edge_str_len - def _update_vertex_degree_list(self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str]) -> None: + def _update_vertex_degree_list( + self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str] + ) -> None: for depth, node_str in enumerate(nodes_with_degree): if depth >= len(vertex_degree_list): vertex_degree_list.append(set()) @@ -395,14 +427,20 @@ def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: schema = self._get_graph_schema() vertex_props_str, edge_props_str = schema.split("\n")[:2] # TODO: rename to vertex (also need update in the schema) - vertex_props_str = vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") - edge_props_str = edge_props_str[len("Edge properties: ") :].strip("[").strip("]") + vertex_props_str = ( + vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") + ) + edge_props_str = ( + edge_props_str[len("Edge properties: ") :].strip("[").strip("]") + ) vertex_labels = self._extract_label_names(vertex_props_str) edge_labels = self._extract_label_names(edge_props_str) return vertex_labels, edge_labels @staticmethod - def _extract_label_names(source: str, head: str = "name: ", tail: str = ", ") -> List[str]: + def _extract_label_names( + source: str, head: str = "name: ", tail: str = ", " + ) -> List[str]: result = [] for s in source.split(head): end = s.find(tail) @@ -421,12 +459,16 @@ def _get_graph_schema(self, refresh: bool = False) -> str: relationships = schema.getRelations() self._schema = ( - f"Vertex properties: {vertex_schema}\nEdge properties: {edge_schema}\nRelationships: {relationships}\n" + f"Vertex properties: {vertex_schema}\n" + f"Edge properties: {edge_schema}\n" + f"Relationships: {relationships}\n" ) log.debug("Link(Relation): %s", relationships) return self._schema - def _limit_property_query(self, value: Optional[str], item_type: str) -> Optional[str]: + def _limit_property_query( + self, value: Optional[str], item_type: str + ) -> Optional[str]: # NOTE: we skip the filter for list/set type (e.g., list of string, add it if needed) if not self._limit_property or not isinstance(value, str): return value 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 d8f59f50e..e9bccc2f4 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 @@ -39,7 +39,11 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} + new_vertex = { + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex + } mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) @@ -47,7 +51,9 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: mini_schema["edgelabels"] = [] for edge in schema["edgelabels"]: new_edge = { - key: edge[key] for key in ["name", "source_label", "target_label", "properties"] if key in edge + key: edge[key] + for key in ["name", "source_label", "target_label", "properties"] + if key in edge } mini_schema["edgelabels"].append(new_edge) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index 4f288a9de..d41f99a37 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -62,10 +62,14 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: # !: We have assumed that self.example is not empty queries = [example["query"] for example in self.examples] # TODO: refactor function chain async to avoid blocking - examples_embedding = asyncio.run(get_embeddings_parallel(self.embedding, queries)) + examples_embedding = asyncio.run( + get_embeddings_parallel(self.embedding, queries) + ) embed_dim = len(examples_embedding[0]) if len(self.examples) > 0: - vector_index = self.vector_index.from_name(embed_dim, self.vector_index_name) + vector_index = self.vector_index.from_name( + embed_dim, self.vector_index_name + ) vector_index.add(examples_embedding, self.examples) <<<<<<< HEAD vector_index.to_index_file(self.index_dir, self.filename_prefix) 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 8021a95b6..ba2b1e379 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 @@ -54,12 +54,18 @@ def __init__(self, embedding: BaseEmbedding): self.vid_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) ======= def __init__(self, embedding: BaseEmbedding, vector_index: type[VectorStoreBase]): +<<<<<<< HEAD self.vid_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids") >>>>>>> 38dce0b (feat(llm): vector db finished) ======= 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") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + self.vid_index = vector_index.from_name( + embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids" + ) +>>>>>>> 3aeef7d (fix) self.embedding = embedding self.sm = SchemaManager(huge_settings.graph_name) @@ -78,12 +84,20 @@ async def get_embeddings_with_semaphore(vid_list: list[str]) -> Any: # This pattern avoids blocking the event loop and prepares for a future fully async pipeline. async with sem: loop = asyncio.get_running_loop() - return await loop.run_in_executor(None, self.embedding.get_texts_embeddings, vid_list) + return await loop.run_in_executor( + None, self.embedding.get_texts_embeddings, vid_list + ) +<<<<<<< HEAD # Split vids into batches of size batch_size vid_batches = [vids[i : i + batch_size] for i in range(0, len(vids), batch_size)] # Create tasks for each batch +======= + vid_batches = [ + vids[i : i + batch_size] for i in range(0, len(vids), batch_size) + ] +>>>>>>> 3aeef7d (fix) tasks = [get_embeddings_with_semaphore(batch) for batch in vid_batches] embeddings = [] @@ -97,22 +111,39 @@ async def get_embeddings_with_semaphore(vid_list: list[str]) -> Any: >>>>>>> 38dce0b (feat(llm): vector db finished) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: vertexlabels = self.sm.schema.getSchema()["vertexlabels"] - all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) + all_pk_flag = all( + data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels + ) past_vids = self.vid_index.get_all_properties() # only support Faiss # TODO: We should build vid vector index separately, especially when the vertices may be very large <<<<<<< HEAD +<<<<<<< HEAD ======= >>>>>>> 38dce0b (feat(llm): vector db finished) present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py +======= + present_vids = context[ + "vertices" + ] # Warning: data truncated by fetch_graph_data.py +>>>>>>> 3aeef7d (fix) removed_vids = set(past_vids) - set(present_vids) removed_num = self.vid_index.remove(removed_vids) added_vids = list(set(present_vids) - set(past_vids)) if added_vids: +<<<<<<< HEAD vids_to_process = self._extract_names(added_vids) if all_pk_flag else added_vids added_embeddings = asyncio.run(get_embeddings_parallel(self.embedding, vids_to_process)) +======= + vids_to_process = ( + self._extract_names(added_vids) if all_pk_flag else added_vids + ) + added_embeddings = asyncio.run( + self._get_embeddings_parallel(vids_to_process) + ) +>>>>>>> 3aeef7d (fix) log.info("Building vector index for %s vertices...", len(added_vids)) self.vid_index.add(added_embeddings, added_vids) <<<<<<< HEAD diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index d76ad021a..8d1f3394f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -42,26 +42,40 @@ def __init__( self.num_examples = num_examples if not vector_index.exist("gremlin_examples"): log.warning("No gremlin example index found, will generate one.") +<<<<<<< HEAD self.vector_index = vector_index.from_name(self.embedding.get_embedding_dim(), "gremlin_examples") <<<<<<< HEAD ======= >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + self.vector_index = vector_index.from_name( + self.embedding.get_embedding_dim(), "gremlin_examples" + ) +>>>>>>> 3aeef7d (fix) self._build_default_example_index() else: - self.vector_index = vector_index.from_name(self.embedding.get_embedding_dim(), "gremlin_examples") + self.vector_index = vector_index.from_name( + self.embedding.get_embedding_dim(), "gremlin_examples" + ) - def _get_match_result(self, context: Dict[str, Any], query: str) -> List[Dict[str, Any]]: + def _get_match_result( + self, context: Dict[str, Any], query: str + ) -> List[Dict[str, Any]]: if self.num_examples <= 0: return [] query_embedding = context.get("query_embedding") if not isinstance(query_embedding, list): query_embedding = self.embedding.get_texts_embeddings([query])[0] - return self.vector_index.search(query_embedding, self.num_examples, dis_threshold=1.8) + return self.vector_index.search( + query_embedding, self.num_examples, dis_threshold=1.8 + ) def _build_default_example_index(self): - properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict(orient="records") + properties = pd.read_csv( + os.path.join(resource_path, "demo", "text2gremlin.csv") + ).to_dict(orient="records") from concurrent.futures import ThreadPoolExecutor # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py index 929799c67..c952e5232 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py @@ -51,6 +51,7 @@ def __init__( vector_dis_threshold: float = huge_settings.vector_dis_threshold, ): <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space @@ -72,6 +73,11 @@ def __init__( >>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) ======= self.index_dir = str(os.path.join(resource_path, huge_settings.graph_name, "graph_vids")) +======= + self.index_dir = str( + os.path.join(resource_path, huge_settings.graph_name, "graph_vids") + ) +>>>>>>> 3aeef7d (fix) self.vector_index = vector_index.from_name( embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids" ) @@ -98,7 +104,9 @@ def _exact_match_vids(self, keywords: List[str]) -> Tuple[List[str], List[str]]: possible_vids.update([f"{i + 1}:{keyword}" for keyword in keywords]) vids_str = ",".join([f"'{vid}'" for vid in possible_vids]) - resp = self._client.gremlin().exec(SemanticIdQuery.ID_QUERY_TEMPL.format(vids_str=vids_str)) + resp = self._client.gremlin().exec( + SemanticIdQuery.ID_QUERY_TEMPL.format(vids_str=vids_str) + ) searched_vids = [v["id"] for v in resp["data"]] unsearched_keywords = set(keywords) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py index 8a2342c1a..97d2e8262 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py @@ -41,6 +41,7 @@ def __init__(self, vector_index: type[VectorStoreBase], embedding: BaseEmbedding self.topk = topk <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space @@ -56,6 +57,11 @@ def __init__(self, vector_index: type[VectorStoreBase], embedding: BaseEmbedding ======= self.vector_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "chunks") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + self.vector_index = vector_index.from_name( + embedding.get_embedding_dim(), huge_settings.graph_name, "chunks" + ) +>>>>>>> 3aeef7d (fix) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query = context.get("query") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py index 0443b15c8..c83fa7781 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py @@ -61,7 +61,9 @@ def __init__(self, llm: BaseLLM, embedding: Optional[BaseEmbedding] = None, grap self.graph = graph self.result = None - def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): + def import_schema( + self, from_hugegraph=None, from_extraction=None, from_user_defined=None + ): if from_hugegraph: self.operators.append(SchemaManager(from_hugegraph)) elif from_user_defined: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py index b5581b256..b0864816d 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py @@ -62,13 +62,23 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: +<<<<<<< HEAD <<<<<<< HEAD context_str = f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") ======= context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip("\n") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + context_str = ( + f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( + "\n" + ) + ) +>>>>>>> 3aeef7d (fix) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) response = self._llm.generate(prompt=final_prompt) return {"answer": response} @@ -84,8 +94,12 @@ def init_llm(self, context): if self._question is None: self._question = context.get("query") or None assert self._question is not None, "No question for synthesizing." - context_head_str = context.get("synthesize_context_head") or self._context_head or "" - context_tail_str = context.get("synthesize_context_tail") or self._context_tail or "" + context_head_str = ( + context.get("synthesize_context_head") or self._context_head or "" + ) + context_tail_str = ( + context.get("synthesize_context_tail") or self._context_tail or "" + ) return context_head_str, context_tail_str def handle_vector_graph(self, context): @@ -98,7 +112,9 @@ def handle_vector_graph(self, context): vector_result_context = "No (vector)phrase related to the query." graph_result = context.get("graph_result") if graph_result: - graph_context_head = context.get("graph_context_head", "Knowledge from graphdb for the query:\n") + graph_context_head = context.get( + "graph_context_head", "Knowledge from graphdb for the query:\n" + ) graph_result_context = graph_context_head + "\n".join( f"{i + 1}. {res}" for i, res in enumerate(graph_result) ) @@ -107,17 +123,29 @@ def handle_vector_graph(self, context): log.warning(graph_result_context) return graph_result_context, vector_result_context - async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[str, Any], None]: + async def run_streaming( + self, context: Dict[str, Any] + ) -> AsyncGenerator[Dict[str, Any], None]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: +<<<<<<< HEAD <<<<<<< HEAD context_str = f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") ======= context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip("\n") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + context_str = ( + f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( + "\n" + ) + ) +>>>>>>> 3aeef7d (fix) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) response = self._llm.generate(prompt=final_prompt) yield {"answer": response} return @@ -141,37 +169,73 @@ async def async_generate( async_tasks = {} if self._raw_answer: final_prompt = self._question - async_tasks["raw_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) + async_tasks["raw_task"] = asyncio.create_task( + self._llm.agenerate(prompt=final_prompt) + ) if self._vector_only_answer: +<<<<<<< HEAD <<<<<<< HEAD context_str = f"{context_head_str}\n" f"{vector_result_context}\n" f"{context_tail_str}".strip("\n") ======= context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip("\n") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip( + "\n" + ) +>>>>>>> 3aeef7d (fix) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) - async_tasks["vector_only_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) + async_tasks["vector_only_task"] = asyncio.create_task( + self._llm.agenerate(prompt=final_prompt) + ) if self._graph_only_answer: +<<<<<<< HEAD <<<<<<< HEAD context_str = f"{context_head_str}\n" f"{graph_result_context}\n" f"{context_tail_str}".strip("\n") ======= context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip("\n") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + context_str = ( + f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( + "\n" + ) + ) +>>>>>>> 3aeef7d (fix) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) - async_tasks["graph_only_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) + async_tasks["graph_only_task"] = asyncio.create_task( + self._llm.agenerate(prompt=final_prompt) + ) if self._graph_vector_answer: context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" +<<<<<<< HEAD <<<<<<< HEAD context_str = f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") ======= context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + context_str = ( + f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip( + "\n" + ) + ) +>>>>>>> 3aeef7d (fix) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) - async_tasks["graph_vector_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) + async_tasks["graph_vector_task"] = asyncio.create_task( + self._llm.agenerate(prompt=final_prompt) + ) async_tasks_mapping = { "raw_task": "raw_answer", @@ -204,17 +268,27 @@ async def async_streaming_generate( if self._raw_answer: final_prompt = self._question async_generators.append( - self.__llm_generate_with_meta_info(task_id=auto_id, target_key="raw_answer", prompt=final_prompt) + self.__llm_generate_with_meta_info( + task_id=auto_id, target_key="raw_answer", prompt=final_prompt + ) ) auto_id += 1 if self._vector_only_answer: +<<<<<<< HEAD <<<<<<< HEAD context_str = f"{context_head_str}\n" f"{vector_result_context}\n" f"{context_tail_str}".strip("\n") ======= context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip("\n") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip( + "\n" + ) +>>>>>>> 3aeef7d (fix) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_generators.append( self.__llm_generate_with_meta_info( task_id=auto_id, target_key="vector_only_answer", prompt=final_prompt @@ -222,28 +296,50 @@ async def async_streaming_generate( ) auto_id += 1 if self._graph_only_answer: +<<<<<<< HEAD <<<<<<< HEAD context_str = f"{context_head_str}\n" f"{graph_result_context}\n" f"{context_tail_str}".strip("\n") ======= context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip("\n") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + context_str = ( + f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( + "\n" + ) + ) +>>>>>>> 3aeef7d (fix) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_generators.append( - self.__llm_generate_with_meta_info(task_id=auto_id, target_key="graph_only_answer", prompt=final_prompt) + self.__llm_generate_with_meta_info( + task_id=auto_id, target_key="graph_only_answer", prompt=final_prompt + ) ) auto_id += 1 if self._graph_vector_answer: context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" +<<<<<<< HEAD <<<<<<< HEAD context_str = f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") ======= context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + context_str = ( + f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip( + "\n" + ) + ) +>>>>>>> 3aeef7d (fix) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_generators.append( self.__llm_generate_with_meta_info( task_id=auto_id, target_key="graph_vector_answer", prompt=final_prompt @@ -256,7 +352,9 @@ async def async_streaming_generate( async_tasks = [asyncio.create_task(anext(gen)) for gen in async_generators] while True: - done, _ = await asyncio.wait(async_tasks, return_when=asyncio.FIRST_COMPLETED) + done, _ = await asyncio.wait( + async_tasks, return_when=asyncio.FIRST_COMPLETED + ) stop_task_num = 0 for task in done: try: @@ -270,7 +368,9 @@ async def async_streaming_generate( break yield context - async def __llm_generate_with_meta_info(self, task_id: int, target_key: str, prompt: str): + async def __llm_generate_with_meta_info( + self, task_id: int, target_key: str, prompt: str + ): # FIXME: Expected type 'AsyncIterable', got 'Coroutine[Any, Any, AsyncGenerator[str, None]]' instead async for token in self._llm.agenerate_streaming(prompt=prompt): yield task_id, target_key, token diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py index 82d985a38..13b23fa2c 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py @@ -52,6 +52,7 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: data["triples"] = [] extract_triples_by_regex(llm_output, data) <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD print( f"LLM {self.__class__.__name__} input:{prompt} \n" @@ -63,6 +64,11 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: ======= print(f"LLM {self.__class__.__name__} input:{prompt} \n output: {llm_output} \n data: {data}") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + print( + f"LLM {self.__class__.__name__} input:{prompt} \n output: {llm_output} \n data: {data}" + ) +>>>>>>> 3aeef7d (fix) data["call_count"] = data.get("call_count", 0) + 1 return data diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index e1d299cc1..2c0244d57 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -48,12 +48,15 @@ def _extract_response(self, response: str, label: str = "gremlin") -> str: return match.group(1).strip() return response.strip() - def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional[str]: + def _format_examples( + self, examples: Optional[List[Dict[str, str]]] + ) -> Optional[str]: if not examples: return None example_strings = [] for example in examples: <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD example_strings.append( f"- query: {example['query']}\n" @@ -65,6 +68,11 @@ def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional ======= example_strings.append(f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + example_strings.append( + f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```" + ) +>>>>>>> 3aeef7d (fix) return "\n\n".join(example_strings) def _format_vertices(self, vertices: Optional[List[str]]) -> Optional[str]: @@ -80,7 +88,9 @@ def _format_properties(self, properties: Optional[List[tuple]]) -> Optional[str] async def async_generate(self, context: Dict[str, Any]): async_tasks = {} query = context.get("query") - raw_example = [{"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"}] + raw_example = [ + {"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"} + ] raw_prompt = self.gremlin_prompt.format( query=query, schema=self.schema, @@ -88,7 +98,9 @@ async def async_generate(self, context: Dict[str, Any]): vertices=self._format_vertices(vertices=self.vertices), properties=self._format_properties(properties=None), ) - async_tasks["raw_answer"] = asyncio.create_task(self.llm.agenerate(prompt=raw_prompt)) + async_tasks["raw_answer"] = asyncio.create_task( + self.llm.agenerate(prompt=raw_prompt) + ) examples = context.get("match_result") init_prompt = self.gremlin_prompt.format( @@ -98,7 +110,9 @@ async def async_generate(self, context: Dict[str, Any]): vertices=self._format_vertices(vertices=self.vertices), properties=self._format_properties(properties=None), ) - async_tasks["initialized_answer"] = asyncio.create_task(self.llm.agenerate(prompt=init_prompt)) + async_tasks["initialized_answer"] = asyncio.create_task( + self.llm.agenerate(prompt=init_prompt) + ) raw_response = await async_tasks["raw_answer"] initialized_response = await async_tasks["initialized_answer"] @@ -116,7 +130,9 @@ async def async_generate(self, context: Dict[str, Any]): def sync_generate(self, context: Dict[str, Any]): query = context.get("query") - raw_example = [{"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"}] + raw_example = [ + {"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"} + ] raw_prompt = self.gremlin_prompt.format( query=query, schema=self.schema, 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 362670a4b..00a237077 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 @@ -75,6 +75,7 @@ def generate_extract_triple_prompt(text, schema=None) -> str: if schema: return schema_real_prompt +<<<<<<< HEAD <<<<<<< HEAD log.warning( <<<<<<< HEAD @@ -87,6 +88,11 @@ def generate_extract_triple_prompt(text, schema=None) -> str: ======= log.warning("Recommend to provide a graph schema to improve the extraction accuracy. Now using the default schema.") >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + log.warning( + "Recommend to provide a graph schema to improve the extraction accuracy. Now using the default schema." + ) +>>>>>>> 3aeef7d (fix) return text_based_prompt @@ -116,6 +122,7 @@ def extract_triples_by_regex_with_schema(schema, text, graph): p_lower = p.lower() for vertex in schema["vertices"]: <<<<<<< HEAD +<<<<<<< HEAD <<<<<<< HEAD if vertex["vertex_label"] == label and any( pp.lower() == p_lower for pp in vertex["properties"] @@ -126,6 +133,11 @@ def extract_triples_by_regex_with_schema(schema, text, graph): ======= if vertex["vertex_label"] == label and any(pp.lower() == p_lower for pp in vertex["properties"]): >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + if vertex["vertex_label"] == label and any( + pp.lower() == p_lower for pp in vertex["properties"] + ): +>>>>>>> 3aeef7d (fix) id = f"{label}-{s}" if id not in vertices_dict: vertices_dict[id] = { @@ -223,6 +235,7 @@ def valid(self, element_id: str, max_length: int = 256) -> bool: return True def _filter_long_id(self, graph) -> Dict[str, List[Any]]: +<<<<<<< HEAD graph["vertices"] = [vertex for vertex in graph["vertices"] if self.valid(vertex["id"])] <<<<<<< HEAD <<<<<<< HEAD @@ -235,4 +248,14 @@ def _filter_long_id(self, graph) -> Dict[str, List[Any]]: ======= graph["edges"] = [edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"])] >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= + graph["vertices"] = [ + vertex for vertex in graph["vertices"] if self.valid(vertex["id"]) + ] + graph["edges"] = [ + edge + for edge in graph["edges"] + if self.valid(edge["start"]) and self.valid(edge["end"]) + ] +>>>>>>> 3aeef7d (fix) return graph diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index 0aa2a2bea..420fc9776 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -64,7 +64,9 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: end_time = time.perf_counter() log.debug("Keyword extraction time: %.2f seconds", end_time - start_time) - keywords = self._extract_keywords_from_response(response=response, lowercase=False, start_token="KEYWORDS:") + keywords = self._extract_keywords_from_response( + response=response, lowercase=False, start_token="KEYWORDS:" + ) keywords = {k.replace("'", "") for k in keywords} context["keywords"] = list(keywords) log.info("User Query: %s\nKeywords: %s", self._query, context["keywords"]) @@ -87,22 +89,32 @@ def _extract_keywords_from_response( match = match[len(start_token) :].strip() <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> 3aeef7d (fix) keywords.extend( k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1 ) +<<<<<<< HEAD ======= keywords.extend(k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1) >>>>>>> 87ee5d3 (style: format code with black line-length 120) ======= keywords.extend(k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1) >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= +>>>>>>> 3aeef7d (fix) # if the keyword consists of multiple words, split into sub-words (removing stopwords) results = set(keywords) for token in keywords: sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: - results.update(w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language)) + results.update( + w + for w in sub_tokens + if w not in NLTKHelper().stopwords(lang=self._language) + ) return results diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py index 058d1bce9..a45812393 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py @@ -31,7 +31,9 @@ def __init__(self, llm: BaseLLM): def _load_few_shot_example(self, example_name: str) -> Dict[str, Any]: """Loads and finds the specified few-shot example from the unified JSON file.""" - examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "prompt_examples.json" + ) if not os.path.exists(examples_path): raise FileNotFoundError(f"Examples file not found: {examples_path}") with open(examples_path, "r", encoding="utf-8") as f: @@ -39,7 +41,9 @@ def _load_few_shot_example(self, example_name: str) -> Dict[str, Any]: for example in all_examples: if example.get("name") == example_name: return example - raise ValueError(f"Example with name '{example_name}' not found in prompt_examples.json") + raise ValueError( + f"Example with name '{example_name}' not found in prompt_examples.json" + ) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: """Executes the core logic of prompt generation.""" @@ -48,7 +52,9 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: example_name = context.get("example_name") if not all([source_text, scenario, example_name]): - raise ValueError("Missing required context: source_text, scenario, or example_name.") + raise ValueError( + "Missing required context: source_text, scenario, or example_name." + ) few_shot_example = self._load_few_shot_example(example_name) meta_prompt = prompt_tpl.generate_extract_prompt_template.format( 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 710760c7d..3ba178c88 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 @@ -67,9 +67,9 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: item_type = item["type"] if item_type == "vertex": label = item["label"] - non_nullable_keys = set(properties_map[item_type][label]["properties"]).difference( - set(properties_map[item_type][label]["nullable_keys"]) - ) + non_nullable_keys = set( + properties_map[item_type][label]["properties"] + ).difference(set(properties_map[item_type][label]["nullable_keys"])) for key in non_nullable_keys: if key not in item["properties"]: item["properties"][key] = "NULL" @@ -82,7 +82,9 @@ 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 + ) -> None: self.llm = llm self.example_prompt = example_prompt self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name @@ -143,8 +145,14 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: try: property_graph = json.loads(json_str) # 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 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'." + ) return items # Create sets for valid vertex and edge labels based on the schema @@ -154,7 +162,9 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: def process_items(item_list, valid_labels, item_type): for item in item_list: if not isinstance(item, dict): - log.warning("Invalid property graph item type '%s'.", type(item)) + log.warning( + "Invalid property graph item type '%s'.", type(item) + ) continue if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): log.warning("Invalid item keys '%s'.", item.keys()) @@ -171,5 +181,7 @@ def process_items(item_list, valid_labels, item_type): process_items(property_graph["vertices"], vertex_label_set, "vertex") process_items(property_graph["edges"], edge_label_set, "edge") except json.JSONDecodeError: - log.critical("Invalid property graph JSON! Please check the extracted JSON data carefully") + log.critical( + "Invalid property graph JSON! Please check the extracted JSON data carefully" + ) return items diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py index 1e33514ca..ae445c206 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py @@ -117,9 +117,13 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: raise ValueError("Context must be a dictionary") if "raw_texts" not in context or not isinstance(context["raw_texts"], list): raise ValueError("'raw_texts' must be a list[str]") - if "query_examples" not in context or not isinstance(context["query_examples"], list): + if "query_examples" not in context or not isinstance( + context["query_examples"], list + ): raise ValueError("'query_examples' must be a list[str]") - if "few_shot_schema" not in context or not isinstance(context["few_shot_schema"], dict): + if "few_shot_schema" not in context or not isinstance( + context["few_shot_schema"], dict + ): raise ValueError("'few_shot_schema' must be a dict") raw_texts = context["raw_texts"] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py index 6beeb0291..98cc97ccf 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py @@ -104,7 +104,9 @@ def nodes_schemas_text_to_list_of_dict(nodes_schemas): properties = json.loads(properties) except json.decoder.JSONDecodeError: properties = {} - result.append({"label": label, "primary_key": primary_key, "properties": properties}) + result.append( + {"label": label, "primary_key": primary_key, "properties": properties} + ) return result @@ -116,7 +118,9 @@ def relationships_schemas_text_to_list_of_dict(relationships_schemas): continue start = relationships_schema_list[0].strip().replace('"', "") end = relationships_schema_list[2].strip().replace('"', "") - relationships_schema_type = relationships_schema_list[1].strip().replace('"', "") + relationships_schema_type = ( + relationships_schema_list[1].strip().replace('"', "") + ) properties = re.search(JSON_REGEX, relationships_schema) if properties is None: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py index 2914c4b28..b5232d268 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py @@ -23,7 +23,9 @@ from hugegraph_llm.utils.log import log -def log_elapsed_time(start_time: float, func: Callable, args: tuple, msg: Optional[str]): +def log_elapsed_time( + start_time: float, func: Callable, args: tuple, msg: Optional[str] +): elapse_time = time.perf_counter() - start_time unit = "s" if elapse_time < 1: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py index 45eb18626..ace3d4b6a 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py @@ -24,13 +24,17 @@ from hugegraph_llm.models.embeddings.base import BaseEmbedding -async def _get_batch_with_progress(embedding: BaseEmbedding, batch: list[str], pbar: tqdm) -> list[Any]: +async def _get_batch_with_progress( + embedding: BaseEmbedding, batch: list[str], pbar: tqdm +) -> list[Any]: result = await embedding.async_get_texts_embeddings(batch) pbar.update(1) return result -async def get_embeddings_parallel(embedding: BaseEmbedding, vids: list[str]) -> list[Any]: +async def get_embeddings_parallel( + embedding: BaseEmbedding, vids: list[str] +) -> list[Any]: """Get embeddings for texts in parallel. This function processes text embeddings asynchronously in parallel, using batching and semaphore @@ -58,7 +62,9 @@ async def get_embeddings_parallel(embedding: BaseEmbedding, vids: list[str]) -> embeddings = [] with tqdm(total=len(vid_batches)) as pbar: # Create tasks for each batch with progress bar updates - tasks = [_get_batch_with_progress(embedding, batch, pbar) for batch in vid_batches] + tasks = [ + _get_batch_with_progress(embedding, batch, pbar) for batch in vid_batches + ] # Use asyncio.gather() to preserve order batch_results = await asyncio.gather(*tasks) @@ -72,7 +78,9 @@ async def get_embeddings_parallel(embedding: BaseEmbedding, vids: list[str]) -> def get_filename_prefix(embedding_type: str = None, model_name: str = None) -> str: """Generate filename based on model name.""" - if not (model_name and model_name.strip() and embedding_type and embedding_type.strip()): + if not ( + model_name and model_name.strip() and embedding_type and embedding_type.strip() + ): return "" # Sanitize model_name to prevent path traversal or invalid filename chars safe_embedding_type = embedding_type.replace("/", "_").replace("\\", "_").strip() diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 8a2130218..6fab74bf9 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -24,21 +24,12 @@ import gradio as gr from hugegraph_llm.flows.scheduler import SchedulerSingleton -<<<<<<< HEAD from .embedding_utils import get_filename_prefix, get_index_folder_name from .hugegraph_utils import get_hg_client, clean_hg_data from .log import log from .vector_index_utils import read_documents -<<<<<<< HEAD from ..config import resource_path, huge_settings, llm_settings from ..indices.vector_index.faiss_vector_store import FaissVectorIndex -======= -from ..config import resource_path, huge_settings -======= -from ..config import huge_settings, index_settings, resource_path ->>>>>>> 38dce0b (feat(llm): vector db finished) -from ..indices.vector_index.faiss_vector_store import FaissVectorIndex ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) from ..models.embeddings.init_embedding import Embeddings from ..models.llms.init_llm import LLMs from ..operators.kg_construction_task import KgBuilder @@ -48,10 +39,10 @@ def get_graph_index_info(): - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) graph_summary_info = builder.fetch_graph_data().run() -<<<<<<< HEAD -<<<<<<< HEAD folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None) @@ -59,16 +50,10 @@ def get_graph_index_info(): vector_index = FaissVectorIndex.from_index_file( str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix, record_miss=False ) -======= - vector_index = get_vector_index_class(index_settings.cur_vector_index) -======= - vector_index = get_vector_index_class(index_settings.cur_vector_index) ->>>>>>> a255aed (fix cycle import & add docs) vector_index_entity = vector_index.from_name( Embeddings().get_embedding().get_embedding_dim(), huge_settings.graph_name, "chunks" ) vector_index_info = vector_index_entity.get_vector_index_info() ->>>>>>> 38dce0b (feat(llm): vector db finished) graph_summary_info["vid_index"] = { "embed_dim": vector_index_info["embed_dim"], "num_vectors": vector_index_info["vector_info"]["chunk_vector_num"], @@ -78,14 +63,6 @@ def get_graph_index_info(): def clean_all_graph_index(): -<<<<<<< HEAD - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None) - ) - FaissVectorIndex.clean(str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix) - FaissVectorIndex.clean(str(os.path.join(resource_path, folder_name, "gremlin_examples")), filename_prefix) -======= # 清理 Faiss 索引目录（兼容默认） faiss_chunks = os.path.join(resource_path, huge_settings.graph_name, "graph_vids") if os.path.isdir(faiss_chunks): @@ -96,7 +73,6 @@ def clean_all_graph_index(): from ..indices.vector_index.faiss_vector_store import FaissVectorIndex FaissVectorIndex.clean("gremlin_examples") ->>>>>>> 87ee5d3 (style: format code with black line-length 120) log.warning("Clear graph index and text2gql index successfully!") gr.Info("Clear graph index and text2gql index successfully!") @@ -166,6 +142,19 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: if not schema: return "ERROR: please input with correct schema/format." + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) + if not schema: + return "ERROR: please input with correct schema/format." + + error_message = parse_schema(schema, builder) + if error_message: + return error_message + builder.chunk_split(texts, "document", "zh").extract_info( + example_prompt, "property_graph" + ) + try: return scheduler.schedule_flow( "graph_extract", schema, texts, example_prompt, "property_graph" @@ -186,8 +175,20 @@ def update_vid_embedding(): def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: - scheduler = SchedulerSingleton.get_instance() - return scheduler.schedule_flow("import_graph_data", data, schema) + data_json = json.loads(data.strip()) + log.debug("Import graph data: %s", data) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) + if schema: + error_message = parse_schema(schema, builder) + if error_message: + return error_message + + context = builder.commit_to_hugegraph().run(data_json) + gr.Info("Import graph data successfully!") + print(context) + return json.dumps(context, ensure_ascii=False, indent=2) except Exception as e: # pylint: disable=W0718 log.error(e) traceback.print_exc() @@ -197,31 +198,7 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: def build_schema(input_text, query_example, few_shot): -<<<<<<< HEAD scheduler = SchedulerSingleton.get_instance() -======= - context = {"raw_texts": [input_text] if input_text else [], "query_examples": [], "few_shot_schema": {}} - - if few_shot: - try: - context["few_shot_schema"] = json.loads(few_shot) - except json.JSONDecodeError as e: - raise gr.Error(f"Few Shot Schema is not in a valid JSON format: {e}") from e - - if query_example: - try: - parsed_examples = json.loads(query_example) - # Validate and retain the description and gremlin fields - context["query_examples"] = [ - {"description": ex.get("description", ""), "gremlin": ex.get("gremlin", "")} - for ex in parsed_examples - if isinstance(ex, dict) and "description" in ex and "gremlin" in ex - ] - except json.JSONDecodeError as e: - raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e - - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) ->>>>>>> 87ee5d3 (style: format code with black line-length 120) try: return scheduler.schedule_flow( "build_schema", input_text, query_example, few_shot diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index bf3b84ecf..082d6bb8c 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -29,7 +29,9 @@ MAX_BACKUP_DIRS = 7 MAX_VERTICES = 100000 MAX_EDGES = 200000 -BACKUP_DIR = str(os.path.join(resource_path, "backup-graph-data-4020", huge_settings.graph_name)) +BACKUP_DIR = str( + os.path.join(resource_path, "backup-graph-data-4020", huge_settings.graph_name) +) def run_gremlin_query(query, fmt=True): @@ -53,22 +55,36 @@ def init_hg_test_data(): schema = client.schema() schema.propertyKey("name").asText().ifNotExist().create() schema.propertyKey("birthDate").asText().ifNotExist().create() - schema.vertexLabel("Person").properties("name", "birthDate").useCustomizeStringId().ifNotExist().create() - schema.vertexLabel("Movie").properties("name").useCustomizeStringId().ifNotExist().create() - schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() - - schema.indexLabel("PersonByName").onV("Person").by("name").secondary().ifNotExist().create() - schema.indexLabel("MovieByName").onV("Movie").by("name").secondary().ifNotExist().create() + schema.vertexLabel("Person").properties( + "name", "birthDate" + ).useCustomizeStringId().ifNotExist().create() + schema.vertexLabel("Movie").properties( + "name" + ).useCustomizeStringId().ifNotExist().create() + schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel( + "Movie" + ).ifNotExist().create() + + schema.indexLabel("PersonByName").onV("Person").by( + "name" + ).secondary().ifNotExist().create() + schema.indexLabel("MovieByName").onV("Movie").by( + "name" + ).secondary().ifNotExist().create() graph = client.graph() - graph.addVertex("Person", {"name": "Al Pacino", "birthDate": "1940-04-25"}, id="Al Pacino") + graph.addVertex( + "Person", {"name": "Al Pacino", "birthDate": "1940-04-25"}, id="Al Pacino" + ) graph.addVertex( "Person", {"name": "Robert De Niro", "birthDate": "1943-08-17"}, id="Robert De Niro", ) graph.addVertex("Movie", {"name": "The Godfather"}, id="The Godfather") - graph.addVertex("Movie", {"name": "The Godfather Part II"}, id="The Godfather Part II") + graph.addVertex( + "Movie", {"name": "The Godfather Part II"}, id="The Godfather Part II" + ) graph.addVertex( "Movie", {"name": "The Godfather Coda The Death of Michael Corleone"}, @@ -77,7 +93,9 @@ def init_hg_test_data(): graph.addEdge("ActedIn", "Al Pacino", "The Godfather", {}) graph.addEdge("ActedIn", "Al Pacino", "The Godfather Part II", {}) - graph.addEdge("ActedIn", "Al Pacino", "The Godfather Coda The Death of Michael Corleone", {}) + graph.addEdge( + "ActedIn", "Al Pacino", "The Godfather Coda The Death of Michael Corleone", {} + ) graph.addEdge("ActedIn", "Robert De Niro", "The Godfather Part II", {}) schema.getSchema() graph.close() @@ -116,7 +134,9 @@ def backup_data(): } vertexlabels = client.schema().getSchema()["vertexlabels"] - all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) + all_pk_flag = all( + data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels + ) for filename, query in files.items(): write_backup_file(client, backup_subdir, filename, query, all_pk_flag) @@ -140,16 +160,22 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): data = ( <<<<<<< HEAD <<<<<<< HEAD +<<<<<<< HEAD +======= +>>>>>>> 3aeef7d (fix) [ {key: value for key, value in vertex.items() if key != "id"} for vertex in data_full ] +<<<<<<< HEAD ======= [{key: value for key, value in vertex.items() if key != "id"} for vertex in data_full] >>>>>>> 87ee5d3 (style: format code with black line-length 120) ======= [{key: value for key, value in vertex.items() if key != "id"} for vertex in data_full] >>>>>>> 8e0bf08 (chore: mark vectordb optional) +======= +>>>>>>> 3aeef7d (fix) if all_pk_flag else data_full ) @@ -158,7 +184,9 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): data_full = query if isinstance(data_full, dict) and "schema" in data_full: groovy_filename = filename.replace(".json", ".groovy") - with open(os.path.join(backup_subdir, groovy_filename), "w", encoding="utf-8") as groovy_file: + with open( + os.path.join(backup_subdir, groovy_filename), "w", encoding="utf-8" + ) as groovy_file: groovy_file.write(str(data_full["schema"])) else: data = data_full @@ -168,7 +196,9 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): def manage_backup_retention(): try: backup_dirs = [ - os.path.join(BACKUP_DIR, d) for d in os.listdir(BACKUP_DIR) if os.path.isdir(os.path.join(BACKUP_DIR, d)) + os.path.join(BACKUP_DIR, d) + for d in os.listdir(BACKUP_DIR) + if os.path.isdir(os.path.join(BACKUP_DIR, d)) ] backup_dirs.sort(key=os.path.getctime) if len(backup_dirs) > MAX_BACKUP_DIRS: @@ -186,7 +216,9 @@ def manage_backup_retention(): # TODO: In the path demo/rag_demo/configs_block.py, # there is a function test_api_connection that is similar to this function, # but it is not straightforward to reuse -def check_graph_db_connection(url: str, name: str, user: str, pwd: str, graph_space: str) -> bool: +def check_graph_db_connection( + url: str, name: str, user: str, pwd: str, graph_space: str +) -> bool: try: if graph_space and graph_space.strip(): test_url = f"{url}/graphspaces/{graph_space}/graphs/{name}/schema" diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index affcf2459..df48ea1e9 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -21,23 +21,12 @@ import docx import gradio as gr -<<<<<<< HEAD -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex -from hugegraph_llm.models.embeddings.init_embedding import Embeddings, model_map -from hugegraph_llm.flows.scheduler import SchedulerSingleton -from hugegraph_llm.utils.embedding_utils import ( - get_filename_prefix, - get_index_folder_name, -) -======= from hugegraph_llm.config import huge_settings, index_settings from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex from hugegraph_llm.indices.vector_index.milvus_vector_store import MilvusVectorIndex from hugegraph_llm.indices.vector_index.qdrant_vector_store import QdrantVectorIndex from hugegraph_llm.models.embeddings.init_embedding import Embeddings ->>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.models.llms.init_llm import LLMs from hugegraph_llm.operators.kg_construction_task import KgBuilder from hugegraph_llm.utils.hugegraph_utils import get_hg_client @@ -62,7 +51,9 @@ def read_documents(input_file, input_text): texts.append(text) elif full_path.endswith(".pdf"): # TODO: support PDF file - raise gr.Error("PDF will be supported later! Try to upload text/docx now") + raise gr.Error( + "PDF will be supported later! Try to upload text/docx now" + ) else: raise gr.Error("Please input txt or docx file.") else: @@ -78,14 +69,10 @@ def get_vector_index_info(): ) return json.dumps( -<<<<<<< HEAD - {**vector_index_entity.get_vector_index_info(), "now_vector_index": index_settings.now_vector_index}, -======= { **vector_index_entity.get_vector_index_info(), "cur_vector_index": index_settings.cur_vector_index, }, ->>>>>>> a255aed (fix cycle import & add docs) ensure_ascii=False, indent=2, ) @@ -102,8 +89,14 @@ def build_vector_index(input_file, input_text): if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - context = builder.chunk_split(texts, "paragraph", "zh").build_vector_index(vector_index).run() + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) + context = ( + builder.chunk_split(texts, "paragraph", "zh") + .build_vector_index(vector_index) + .run() + ) return json.dumps(context, ensure_ascii=False, indent=2) From f6fa0b7fc98bfc5061a7cfbe840be92d64f5aa6c Mon Sep 17 00:00:00 2001 From: lingxiao Date: Thu, 28 Aug 2025 17:45:57 +0800 Subject: [PATCH 26/71] fix --- .../hugegraph_llm/api/models/rag_requests.py | 24 +- .../src/hugegraph_llm/api/rag_api.py | 23 +- .../src/hugegraph_llm/config/llm_config.py | 16 +- .../config/models/base_config.py | 12 +- .../config/models/base_prompt_config.py | 12 +- .../demo/rag_demo/admin_block.py | 4 +- .../src/hugegraph_llm/demo/rag_demo/app.py | 4 +- .../demo/rag_demo/configs_block.py | 554 ++---------------- .../demo/rag_demo/other_block.py | 12 +- .../hugegraph_llm/demo/rag_demo/rag_block.py | 165 ++---- .../demo/rag_demo/text2gremlin_block.py | 160 +---- .../demo/rag_demo/vector_graph_block.py | 38 +- .../vector_index/milvus_vector_store.py | 8 +- .../vector_index/qdrant_vector_store.py | 8 +- .../hugegraph_llm/middleware/middleware.py | 4 +- .../models/embeddings/init_embedding.py | 8 +- .../hugegraph_llm/models/embeddings/openai.py | 4 +- .../src/hugegraph_llm/models/llms/init_llm.py | 6 +- .../src/hugegraph_llm/models/llms/ollama.py | 8 +- .../src/hugegraph_llm/models/llms/openai.py | 12 +- .../hugegraph_llm/models/rerankers/cohere.py | 4 +- .../models/rerankers/siliconflow.py | 4 +- .../operators/common_op/check_schema.py | 40 +- .../operators/common_op/merge_dedup_rerank.py | 10 +- .../operators/common_op/nltk_helper.py | 4 +- .../operators/document_op/word_extract.py | 6 +- .../hugegraph_llm/operators/graph_rag_task.py | 4 +- .../operators/gremlin_generate_task.py | 16 +- .../hugegraph_op/commit_to_hugegraph.py | 62 +- .../hugegraph_op/fetch_graph_data.py | 4 +- .../operators/hugegraph_op/graph_rag_query.py | 44 +- .../operators/hugegraph_op/schema_manager.py | 4 +- .../index_op/build_gremlin_example_index.py | 8 +- .../index_op/build_semantic_index.py | 85 +-- .../index_op/gremlin_example_index_query.py | 14 +- .../operators/index_op/semantic_id_query.py | 53 +- .../operators/kg_construction_task.py | 4 +- .../operators/llm_op/answer_synthesize.py | 160 ++--- .../operators/llm_op/gremlin_generate.py | 16 +- .../operators/llm_op/info_extract.py | 65 +- .../operators/llm_op/keyword_extract.py | 4 +- .../operators/llm_op/prompt_generate.py | 12 +- .../llm_op/property_graph_extract.py | 18 +- .../operators/llm_op/schema_build.py | 8 +- .../llm_op/unstructured_data_utils.py | 8 +- .../src/hugegraph_llm/utils/decorators.py | 4 +- .../hugegraph_llm/utils/embedding_utils.py | 12 +- .../hugegraph_llm/utils/graph_index_utils.py | 118 ++-- .../hugegraph_llm/utils/hugegraph_utils.py | 40 +- .../hugegraph_llm/utils/vector_index_utils.py | 14 +- 50 files changed, 370 insertions(+), 1557 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py index 9ebc94eb1..f46aea02c 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py +++ b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py @@ -34,18 +34,12 @@ class GraphConfigRequest(BaseModel): class RAGRequest(BaseModel): query: str = Query(..., description="Query you want to ask") raw_answer: bool = Query(False, description="Use LLM to generate answer directly") - vector_only: bool = Query( - False, description="Use LLM to generate answer with vector" - ) - graph_only: bool = Query( - True, description="Use LLM to generate answer with graph RAG only" - ) + vector_only: bool = Query(False, description="Use LLM to generate answer with vector") + graph_only: bool = Query(True, description="Use LLM to generate answer with graph RAG only") graph_vector_answer: bool = Query( False, description="Use LLM to generate answer with vector & GraphRAG" ) - graph_ratio: float = Query( - 0.5, description="The ratio of GraphRAG ans & vector ans" - ) + graph_ratio: float = Query(0.5, description="The ratio of GraphRAG ans & vector ans") rerank_method: Literal["bleu", "reranker"] = Query( "bleu", description="Method to rerank the results." ) @@ -59,9 +53,7 @@ class RAGRequest(BaseModel): max_graph_items: int = Query( 30, description="Maximum number of items for GQL queries in graph." ) - topk_return_results: int = Query( - 20, description="Number of sorted results to return finally." - ) + topk_return_results: int = Query(20, description="Number of sorted results to return finally.") vector_dis_threshold: float = Query( 0.9, description="Threshold for vector similarity\ @@ -98,9 +90,7 @@ class GraphRAGRequest(BaseModel): max_graph_items: int = Query( 30, description="Maximum number of items for GQL queries in graph." ) - topk_return_results: int = Query( - 20, description="Number of sorted results to return finally." - ) + topk_return_results: int = Query(20, description="Number of sorted results to return finally.") vector_dis_threshold: float = Query( 0.9, description="Threshold for vector similarity \ @@ -115,9 +105,7 @@ class GraphRAGRequest(BaseModel): client_config: Optional[GraphConfigRequest] = Query( None, description="hugegraph server config." ) - get_vertex_only: bool = Query( - False, description="return only keywords & vertex (early stop)." - ) + get_vertex_only: bool = Query(False, description="return only keywords & vertex (early stop).") gremlin_tmpl_num: int = Query( 1, diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py index 1d5b451b1..5c9295efa 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py @@ -67,8 +67,7 @@ def rag_answer_api(req: RAGRequest): # Keep prompt params in the end custom_related_information=req.custom_priority_info, answer_prompt=req.answer_prompt or prompt.answer_prompt, - keywords_extract_prompt=req.keywords_extract_prompt - or prompt.keywords_extract_prompt, + keywords_extract_prompt=req.keywords_extract_prompt or prompt.keywords_extract_prompt, gremlin_prompt=req.gremlin_prompt or prompt.gremlin_generate_prompt, ) # TODO: we need more info in the response for users to understand the query logic @@ -136,9 +135,7 @@ def graph_rag_recall_api(req: GraphRAGRequest): except TypeError as e: log.error("TypeError in graph_rag_recall_api: %s", e) - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) - ) from e + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e except Exception as e: log.error("Unexpected error occurred: %s", e) raise HTTPException( @@ -149,9 +146,7 @@ def graph_rag_recall_api(req: GraphRAGRequest): @router.post("/config/graph", status_code=status.HTTP_201_CREATED) def graph_config_api(req: GraphConfigRequest): # Accept status code - res = apply_graph_conf( - req.url, req.name, req.user, req.pwd, req.gs, origin_call="http" - ) + res = apply_graph_conf(req.url, req.name, req.user, req.pwd, req.gs, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) # TODO: restructure the implement of llm to three types, like "/config/chat_llm" @@ -164,9 +159,7 @@ def llm_config_api(req: LLMConfigRequest): req.api_key, req.api_base, req.language_model, req.max_tokens, origin_call="http" ) else: - res = apply_llm_conf( - req.host, req.port, req.language_model, None, origin_call="http" - ) + res = apply_llm_conf(req.host, req.port, req.language_model, None, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/embedding", status_code=status.HTTP_201_CREATED) @@ -178,9 +171,7 @@ def embedding_config_api(req: LLMConfigRequest): req.api_key, req.api_base, req.language_model, origin_call="http" ) else: - res = apply_embedding_conf( - req.host, req.port, req.language_model, origin_call="http" - ) + res = apply_embedding_conf(req.host, req.port, req.language_model, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/rerank", status_code=status.HTTP_201_CREATED) @@ -192,9 +183,7 @@ def rerank_config_api(req: RerankerConfigRequest): req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http" ) elif req.reranker_type == "siliconflow": - res = apply_reranker_conf( - req.api_key, req.reranker_model, None, origin_call="http" - ) + res = apply_reranker_conf(req.api_key, req.reranker_model, None, origin_call="http") else: res = status.HTTP_501_NOT_IMPLEMENTED return generate_response(RAGResponse(status_code=res, message="Missing Value")) diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index 8b4f274ee..af22b71b0 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -32,19 +32,13 @@ class LLMConfig(BaseConfig): embedding_type: Optional[Literal["openai", "litellm", "ollama/local"]] = "openai" reranker_type: Optional[Literal["cohere", "siliconflow"]] = None # 1. OpenAI settings - openai_chat_api_base: str = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) + openai_chat_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") openai_chat_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_chat_language_model: str = "gpt-4.1-mini" - openai_extract_api_base: str = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) + openai_extract_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") openai_extract_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_extract_language_model: str = "gpt-4.1-mini" - openai_text2gql_api_base: str = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) + openai_text2gql_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") openai_text2gql_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_text2gql_language_model: str = "gpt-4.1-mini" openai_embedding_api_base: str = os.environ.get( @@ -57,9 +51,7 @@ class LLMConfig(BaseConfig): openai_extract_tokens: int = 256 openai_text2gql_tokens: int = 4096 # 2. Rerank settings - cohere_base_url: str = os.environ.get( - "CO_API_URL", "https://api.cohere.com/v1/rerank" - ) + cohere_base_url: str = os.environ.get("CO_API_URL", "https://api.cohere.com/v1/rerank") reranker_api_key: Optional[str] = None reranker_model: Optional[str] = None # 3. Ollama settings diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py index 4ec9256c5..5fec3a778 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py @@ -24,9 +24,7 @@ from hugegraph_llm.utils.log import log dir_name = os.path.dirname -env_path = os.path.join( - os.getcwd(), ".env" -) # Load .env from the current working directory +env_path = os.path.join(os.getcwd(), ".env") # Load .env from the current working directory class BaseConfig(BaseSettings): @@ -90,9 +88,7 @@ def check_env(self): # Step 2: Add missing config items to .env self._sync_object_to_env(env_config, config_dict) except Exception as e: - log.error( - "An error occurred when checking the .env variable file: %s", str(e) - ) + log.error("An error occurred when checking the .env variable file: %s", str(e)) raise def _sync_env_to_object(self, env_config, config_dict): @@ -143,9 +139,7 @@ def __init__(self, **data): # Synchronize configurations between the object and .env file self.check_env() - log.info( - "The %s file was loaded. Class: %s", env_path, self.__class__.__name__ - ) + log.info("The %s file was loaded. Class: %s", env_path, self.__class__.__name__) except Exception as e: log.error("An error occurred when initializing the configuration object: %s", str(e)) raise diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py index b15bad0a6..2369d01a6 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py @@ -57,9 +57,7 @@ def ensure_yaml_file_exists(self): current_dir = Path.cwd().resolve() project_root = get_project_root() if current_dir == project_root: - log.info( - "Current working directory is the project root, proceeding to run the app." - ) + log.info("Current working directory is the project root, proceeding to run the app.") else: error_msg = ( f"Current working directory is not the project root. " @@ -124,9 +122,7 @@ def to_literal(val): "gremlin_generate_prompt": to_literal(self.gremlin_generate_prompt), "doc_input_text": to_literal(self.doc_input_text), "_language_generated": str(self.llm_settings.language).lower().strip(), - "generate_extract_prompt_template": to_literal( - self.generate_extract_prompt_template - ), + "generate_extract_prompt_template": to_literal(self.generate_extract_prompt_template), } with open(yaml_file_path, "w", encoding="utf-8") as file: yaml.dump(data, file, allow_unicode=True, sort_keys=False, default_flow_style=False) @@ -154,9 +150,7 @@ def generate_yaml_file(self): self.keywords_extract_prompt = self.keywords_extract_prompt_EN self.doc_input_text = self.doc_input_text_EN self.save_to_yaml() - log.info( - "Prompt file '%s' has been generated with default values.", yaml_file_path - ) + log.info("Prompt file '%s' has been generated with default values.", yaml_file_path) def update_yaml_file(self): self.save_to_yaml() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index d3beebcbb..1a157d387 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -143,9 +143,7 @@ def create_admin_block(): with gr.Row(): with gr.Column(): # Button to clear LLM Server log, initially hidden - clear_llm_server_button = gr.Button( - "Clear LLM Server Log", visible=False - ) + clear_llm_server_button = gr.Button("Clear LLM Server Log", visible=False) with gr.Column(): # Button to refresh LLM Server log manually refresh_llm_server_button = gr.Button( 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 a78e62361..b0979763c 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -166,9 +166,7 @@ def create_app(): prompt.update_yaml_file() assert admin_settings.enable_login auth_enabled = admin_settings.enable_login.lower() == "true" - log.info( - "(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled" - ) + log.info("(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled") api_auth = APIRouter(dependencies=[Depends(authenticate)] if auth_enabled else []) hugegraph_llm = init_rag_ui() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 155a4b0b8..d4415700b 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -35,7 +35,6 @@ def test_litellm_embedding(api_key, api_base, model_name, model_dim) -> int: llm_client = LiteLLMEmbedding( - embedding_dimension=model_dim, api_key=api_key, api_base=api_base, model_name=model_name, @@ -72,9 +71,7 @@ def test_api_connection( log.debug("Request URL: %s", url) try: if method.upper() == "GET": - resp = requests.get( - url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth - ) + resp = requests.get(url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth) elif method.upper() == "POST": resp = requests.post( url, @@ -109,16 +106,7 @@ def test_api_connection( return resp.status_code -<<<<<<< HEAD -<<<<<<< HEAD -def apply_embedding_config(arg1, arg2, arg3, origin_call=None) -> int: -======= def config_qianfan_model(arg1, arg2, arg3=None, settings_prefix=None, origin_call=None) -> int: -======= -def config_qianfan_model( - arg1, arg2, arg3=None, settings_prefix=None, origin_call=None -) -> int: ->>>>>>> 3aeef7d (fix) setattr(llm_settings, f"qianfan_{settings_prefix}_api_key", arg1) setattr(llm_settings, f"qianfan_{settings_prefix}_secret_key", arg2) if arg3: @@ -129,13 +117,15 @@ def config_qianfan_model( "client_secret": arg2, } status_code = test_api_connection( - "https://aip.baidubce.com/oauth/2.0/token", "POST", params=params, origin_call=origin_call + "https://aip.baidubce.com/oauth/2.0/token", + "POST", + params=params, + origin_call=origin_call, ) return status_code def apply_embedding_config(arg1, arg2, arg3, arg4, origin_call=None) -> int: ->>>>>>> 38dce0b (feat(llm): vector db finished) status_code = -1 embedding_option = llm_settings.embedding_type arg4 = int(arg4) @@ -147,35 +137,15 @@ def apply_embedding_config(arg1, arg2, arg3, arg4, origin_call=None) -> int: test_url = llm_settings.openai_embedding_api_base + "/embeddings" headers = {"Authorization": f"Bearer {arg1}"} data = {"model": arg3, "input": "test"} -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - status_code = test_api_connection( - test_url, method="POST", headers=headers, body=data, origin_call=origin_call - ) -======= - status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) - elif embedding_option == "qianfan_wenxin": - status_code = config_qianfan_model(arg1, arg2, settings_prefix="embedding", origin_call=origin_call) - llm_settings.qianfan_embedding_model = arg3 - llm_settings.qianfan_embedding_model_dim = arg4 ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= status_code = test_api_connection( test_url, method="POST", headers=headers, body=data, origin_call=origin_call ) ->>>>>>> 3aeef7d (fix) elif embedding_option == "ollama/local": llm_settings.ollama_embedding_host = arg1 llm_settings.ollama_embedding_port = int(arg2) llm_settings.ollama_embedding_model = arg3 llm_settings.ollama_embedding_model_dim = arg4 - status_code = test_api_connection( - f"http://{arg1}:{arg2}", origin_call=origin_call - ) + status_code = test_api_connection(f"http://{arg1}:{arg2}", origin_call=origin_call) elif embedding_option == "litellm": llm_settings.litellm_embedding_api_key = arg1 llm_settings.litellm_embedding_api_base = arg2 @@ -266,63 +236,22 @@ def apply_llm_config( setattr(llm_settings, f"openai_{current_llm_config}_language_model", model_name) setattr(llm_settings, f"openai_{current_llm_config}_tokens", int(max_tokens)) -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD test_url = ( getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" ) -======= - test_url = getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" ->>>>>>> 87ee5d3 (style: format code with black line-length 120) -======= - test_url = getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= - test_url = ( - getattr(llm_settings, f"openai_{current_llm_config}_api_base") - + "/chat/completions" - ) ->>>>>>> 3aeef7d (fix) data = { "model": model_name, "temperature": 0.01, "messages": [{"role": "user", "content": "test"}], } -<<<<<<< HEAD headers = {"Authorization": f"Bearer {api_key_or_host}"} -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD status_code = test_api_connection( test_url, method="POST", headers=headers, body=data, origin_call=origin_call ) -======= - headers = {"Authorization": f"Bearer {arg1}"} - status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) -======= - status_code = test_api_connection( - test_url, method="POST", headers=headers, body=data, origin_call=origin_call - ) ->>>>>>> 3aeef7d (fix) - - elif llm_option == "qianfan_wenxin": - status_code = config_qianfan_model( - arg1, arg2, arg3, settings_prefix=current_llm_config, origin_call=origin_call - ) # pylint: disable=C0301 ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) ->>>>>>> 87ee5d3 (style: format code with black line-length 120) -======= - status_code = test_api_connection(test_url, method="POST", headers=headers, body=data, origin_call=origin_call) ->>>>>>> 8e0bf08 (chore: mark vectordb optional) elif llm_option == "ollama/local": setattr(llm_settings, f"ollama_{current_llm_config}_host", api_key_or_host) - setattr( - llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port) - ) + setattr(llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port)) setattr(llm_settings, f"ollama_{current_llm_config}_language_model", model_name) status_code = test_api_connection( f"http://{api_key_or_host}:{api_base_or_port}", origin_call=origin_call @@ -330,12 +259,8 @@ def apply_llm_config( elif llm_option == "litellm": setattr(llm_settings, f"litellm_{current_llm_config}_api_key", api_key_or_host) - setattr( - llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port - ) - setattr( - llm_settings, f"litellm_{current_llm_config}_language_model", model_name - ) + setattr(llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port) + setattr(llm_settings, f"litellm_{current_llm_config}_language_model", model_name) setattr(llm_settings, f"litellm_{current_llm_config}_tokens", int(max_tokens)) status_code = test_litellm_chat( @@ -360,11 +285,6 @@ def create_configs_block() -> list: info="IP:PORT (e.g. 127.0.0.1:8080) or full URL (e.g. http://127.0.0.1:8080)", ), gr.Textbox( -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -======= ->>>>>>> a255aed (fix cycle import & add docs) value=huge_settings.graph_name, label="graph", info="The graph name of HugeGraph-Server instance", @@ -376,31 +296,9 @@ def create_configs_block() -> list: ), gr.Textbox( value=huge_settings.graph_pwd, -<<<<<<< HEAD label="pwd", type="password", info="Password for graph server auth", -======= - value=huge_settings.graph_name, label="graph", info="The graph name of HugeGraph-Server instance" -======= - value=lambda: huge_settings.graph_name, - label="graph", - info="The graph name of HugeGraph-Server instance", ->>>>>>> f42fa9b (feat(llm): use lambda) - ), - gr.Textbox(value=lambda: huge_settings.graph_user, label="user", info="Username for graph server auth"), - gr.Textbox( -<<<<<<< HEAD - value=huge_settings.graph_pwd, label="pwd", type="password", info="Password for graph server auth" ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - value=lambda: huge_settings.graph_pwd, -======= ->>>>>>> a255aed (fix cycle import & add docs) - label="pwd", - type="password", - info="Password for graph server auth", ->>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( value=huge_settings.graph_space, @@ -419,15 +317,9 @@ def create_configs_block() -> list: "> Tips: The OpenAI option also support openai style api from other providers. " "**Refresh the page** to load the **latest configs** in __UI__." ) -<<<<<<< HEAD with gr.Tab(label="chat"): chat_llm_dropdown = gr.Dropdown( choices=["openai", "litellm", "ollama/local"], -======= - with gr.Tab(label='chat'): - chat_llm_dropdown = gr.Dropdown( - choices=["openai", "litellm", "qianfan_wenxin", "ollama/local"], ->>>>>>> 38dce0b (feat(llm): vector db finished) value=getattr(llm_settings, "chat_llm_type"), label="type", ) @@ -439,11 +331,6 @@ def chat_llm_settings(llm_type): if llm_type == "openai": llm_config_input = [ gr.Textbox( -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -======= ->>>>>>> a255aed (fix cycle import & add docs) value=getattr(llm_settings, "openai_chat_api_key"), label="api_key", type="password", @@ -475,69 +362,14 @@ def chat_llm_settings(llm_type): value=getattr(llm_settings, "ollama_chat_language_model"), label="model_name", ), -======= - value=getattr(llm_settings, "openai_chat_api_key"), label="api_key", type="password" -======= - value=lambda: getattr(llm_settings, "openai_chat_api_key"), label="api_key", type="password" ->>>>>>> f42fa9b (feat(llm): use lambda) - ), - gr.Textbox(value=lambda: getattr(llm_settings, "openai_chat_api_base"), label="api_base"), - gr.Textbox( - value=lambda: getattr(llm_settings, "openai_chat_language_model"), label="model_name" - ), - gr.Textbox(value=lambda: getattr(llm_settings, "openai_chat_tokens"), label="max_token"), - ] - elif llm_type == "ollama/local": - llm_config_input = [ - gr.Textbox(value=lambda: getattr(llm_settings, "ollama_chat_host"), label="host"), - gr.Textbox(value=lambda: str(getattr(llm_settings, "ollama_chat_port")), label="port"), - gr.Textbox( - value=lambda: getattr(llm_settings, "ollama_chat_language_model"), label="model_name" - ), - gr.Textbox(value="", visible=False), - ] - elif llm_type == "qianfan_wenxin": - llm_config_input = [ - gr.Textbox( - value=lambda: getattr(llm_settings, "qianfan_chat_api_key"), - label="api_key", - type="password", - ), - gr.Textbox( - value=lambda: getattr(llm_settings, "qianfan_chat_secret_key"), - label="secret_key", - type="password", - ), - gr.Textbox( - value=lambda: getattr(llm_settings, "qianfan_chat_language_model"), label="model_name" - ), -<<<<<<< HEAD - gr.Textbox(value=getattr(llm_settings, "qianfan_chat_language_model"), label="model_name"), ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= ->>>>>>> f42fa9b (feat(llm): use lambda) gr.Textbox(value="", visible=False), ] elif llm_type == "litellm": llm_config_input = [ gr.Textbox( -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD value=getattr(llm_settings, "litellm_chat_api_key"), label="api_key", type="password", -======= - value=getattr(llm_settings, "litellm_chat_api_key"), label="api_key", type="password" ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - value=lambda: getattr(llm_settings, "litellm_chat_api_key"), -======= - value=getattr(llm_settings, "litellm_chat_api_key"), ->>>>>>> a255aed (fix cycle import & add docs) - label="api_key", - type="password", ->>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( value=getattr(llm_settings, "litellm_chat_api_base"), @@ -549,27 +381,15 @@ def chat_llm_settings(llm_type): label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), -<<<<<<< HEAD -<<<<<<< HEAD gr.Textbox( value=getattr(llm_settings, "litellm_chat_tokens"), label="max_token", ), -======= - gr.Textbox(value=getattr(llm_settings, "litellm_chat_tokens"), label="max_token"), ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - gr.Textbox(value=lambda: getattr(llm_settings, "litellm_chat_tokens"), label="max_token"), ->>>>>>> f42fa9b (feat(llm): use lambda) ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_chat_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_chat_op, inputs=llm_config_input) # Determine whether there are Settings in the.env file env_path = os.path.join( os.getcwd(), ".env" @@ -582,31 +402,13 @@ def chat_llm_settings(llm_type): apply_llm_config_with_text2gql_op, inputs=llm_config_input ) if not api_text2sql_key: -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - llm_config_button.click( - apply_llm_config_with_extract_op, inputs=llm_config_input - ) -======= - llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= llm_config_button.click( apply_llm_config_with_extract_op, inputs=llm_config_input ) ->>>>>>> 3aeef7d (fix) with gr.Tab(label="mini_tasks"): extract_llm_dropdown = gr.Dropdown( choices=["openai", "litellm", "ollama/local"], -======= - llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) - - with gr.Tab(label="mini_tasks"): - extract_llm_dropdown = gr.Dropdown( - choices=["openai", "litellm", "qianfan_wenxin", "ollama/local"], ->>>>>>> 38dce0b (feat(llm): vector db finished) value=getattr(llm_settings, "extract_llm_type"), label="type", ) @@ -618,11 +420,6 @@ def extract_llm_settings(llm_type): if llm_type == "openai": llm_config_input = [ gr.Textbox( -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -======= ->>>>>>> a255aed (fix cycle import & add docs) value=getattr(llm_settings, "openai_extract_api_key"), label="api_key", type="password", @@ -632,9 +429,7 @@ def extract_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr( - llm_settings, "openai_extract_language_model" - ), + value=getattr(llm_settings, "openai_extract_language_model"), label="model_name", ), gr.Textbox( @@ -653,76 +448,17 @@ def extract_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr( - llm_settings, "ollama_extract_language_model" - ), + value=getattr(llm_settings, "ollama_extract_language_model"), label="model_name", ), -======= - value=getattr(llm_settings, "openai_extract_api_key"), label="api_key", type="password" -======= - value=lambda: getattr(llm_settings, "openai_extract_api_key"), - label="api_key", - type="password", ->>>>>>> f42fa9b (feat(llm): use lambda) - ), - gr.Textbox(value=lambda: getattr(llm_settings, "openai_extract_api_base"), label="api_base"), - gr.Textbox( - value=lambda: getattr(llm_settings, "openai_extract_language_model"), label="model_name" - ), - gr.Textbox(value=lambda: getattr(llm_settings, "openai_extract_tokens"), label="max_token"), - ] - elif llm_type == "ollama/local": - llm_config_input = [ - gr.Textbox(value=lambda: getattr(llm_settings, "ollama_extract_host"), label="host"), - gr.Textbox(value=lambda: str(getattr(llm_settings, "ollama_extract_port")), label="port"), - gr.Textbox( - value=lambda: getattr(llm_settings, "ollama_extract_language_model"), label="model_name" - ), - gr.Textbox(value="", visible=False), - ] - elif llm_type == "qianfan_wenxin": - llm_config_input = [ - gr.Textbox( - value=lambda: getattr(llm_settings, "qianfan_extract_api_key"), - label="api_key", - type="password", - ), - gr.Textbox( - value=lambda: getattr(llm_settings, "qianfan_extract_secret_key"), - label="secret_key", - type="password", - ), -<<<<<<< HEAD - gr.Textbox(value=getattr(llm_settings, "qianfan_extract_language_model"), label="model_name"), ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - gr.Textbox( - value=lambda: getattr(llm_settings, "qianfan_extract_language_model"), label="model_name" - ), ->>>>>>> f42fa9b (feat(llm): use lambda) gr.Textbox(value="", visible=False), ] elif llm_type == "litellm": llm_config_input = [ gr.Textbox( -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - value=getattr(llm_settings, "litellm_extract_api_key"), - label="api_key", - type="password", -======= - value=getattr(llm_settings, "litellm_extract_api_key"), label="api_key", type="password" ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - value=lambda: getattr(llm_settings, "litellm_extract_api_key"), -======= value=getattr(llm_settings, "litellm_extract_api_key"), ->>>>>>> a255aed (fix cycle import & add docs) label="api_key", type="password", ->>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( value=getattr(llm_settings, "litellm_extract_api_base"), @@ -730,46 +466,23 @@ def extract_llm_settings(llm_type): info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr( - llm_settings, "litellm_extract_language_model" - ), + value=getattr(llm_settings, "litellm_extract_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), -<<<<<<< HEAD -<<<<<<< HEAD gr.Textbox( value=getattr(llm_settings, "litellm_extract_tokens"), label="max_token", ), -======= - gr.Textbox(value=getattr(llm_settings, "litellm_extract_tokens"), label="max_token"), ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - gr.Textbox(value=lambda: getattr(llm_settings, "litellm_extract_tokens"), label="max_token"), ->>>>>>> f42fa9b (feat(llm): use lambda) ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_extract_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) -<<<<<<< HEAD -<<<<<<< HEAD -======= ->>>>>>> 87ee5d3 (style: format code with black line-length 120) with gr.Tab(label="text2gql"): text2gql_llm_dropdown = gr.Dropdown( choices=["openai", "litellm", "ollama/local"], -======= - with gr.Tab(label='text2gql'): - text2gql_llm_dropdown = gr.Dropdown( - choices=["openai", "litellm", "qianfan_wenxin", "ollama/local"], ->>>>>>> 38dce0b (feat(llm): vector db finished) value=getattr(llm_settings, "text2gql_llm_type"), label="type", ) @@ -781,11 +494,6 @@ def text2gql_llm_settings(llm_type): if llm_type == "openai": llm_config_input = [ gr.Textbox( -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -======= ->>>>>>> a255aed (fix cycle import & add docs) value=getattr(llm_settings, "openai_text2gql_api_key"), label="api_key", type="password", @@ -795,9 +503,7 @@ def text2gql_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr( - llm_settings, "openai_text2gql_language_model" - ), + value=getattr(llm_settings, "openai_text2gql_language_model"), label="model_name", ), gr.Textbox( @@ -816,128 +522,41 @@ def text2gql_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr( - llm_settings, "ollama_text2gql_language_model" - ), + value=getattr(llm_settings, "ollama_text2gql_language_model"), label="model_name", ), -======= - value=getattr(llm_settings, "openai_text2gql_api_key"), label="api_key", type="password" -======= - value=lambda: getattr(llm_settings, "openai_text2gql_api_key"), - label="api_key", - type="password", ->>>>>>> f42fa9b (feat(llm): use lambda) - ), - gr.Textbox(value=lambda: getattr(llm_settings, "openai_text2gql_api_base"), label="api_base"), - gr.Textbox( - value=lambda: getattr(llm_settings, "openai_text2gql_language_model"), label="model_name" - ), - gr.Textbox(value=lambda: getattr(llm_settings, "openai_text2gql_tokens"), label="max_token"), - ] - elif llm_type == "ollama/local": - llm_config_input = [ - gr.Textbox(value=lambda: getattr(llm_settings, "ollama_text2gql_host"), label="host"), - gr.Textbox(value=lambda: str(getattr(llm_settings, "ollama_text2gql_port")), label="port"), - gr.Textbox( - value=lambda: getattr(llm_settings, "ollama_text2gql_language_model"), label="model_name" - ), - gr.Textbox(value="", visible=False), - ] - elif llm_type == "qianfan_wenxin": - llm_config_input = [ - gr.Textbox( - value=lambda: getattr(llm_settings, "qianfan_text2gql_api_key"), - label="api_key", - type="password", - ), - gr.Textbox( - value=lambda: getattr(llm_settings, "qianfan_text2gql_secret_key"), - label="secret_key", - type="password", - ), -<<<<<<< HEAD - gr.Textbox(value=getattr(llm_settings, "qianfan_text2gql_language_model"), label="model_name"), ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - gr.Textbox( - value=lambda: getattr(llm_settings, "qianfan_text2gql_language_model"), label="model_name" - ), ->>>>>>> f42fa9b (feat(llm): use lambda) gr.Textbox(value="", visible=False), ] elif llm_type == "litellm": llm_config_input = [ gr.Textbox( -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - value=getattr(llm_settings, "litellm_text2gql_api_key"), - label="api_key", - type="password", -======= - value=getattr(llm_settings, "litellm_text2gql_api_key"), label="api_key", type="password" ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= value=lambda: getattr(llm_settings, "litellm_text2gql_api_key"), -======= - value=lambda: getattr( - llm_settings, "litellm_text2gql_api_key" - ), ->>>>>>> 3aeef7d (fix) label="api_key", type="password", ->>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=lambda: getattr( - llm_settings, "litellm_text2gql_api_base" - ), + value=lambda: getattr(llm_settings, "litellm_text2gql_api_base"), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=lambda: getattr( - llm_settings, "litellm_text2gql_language_model" - ), + value=lambda: getattr(llm_settings, "litellm_text2gql_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), -<<<<<<< HEAD -<<<<<<< HEAD gr.Textbox( -<<<<<<< HEAD - value=getattr(llm_settings, "litellm_text2gql_tokens"), -======= - value=lambda: getattr( - llm_settings, "litellm_text2gql_tokens" - ), ->>>>>>> 3aeef7d (fix) + value=lambda: getattr(llm_settings, "litellm_text2gql_tokens"), label="max_token", ), -======= - gr.Textbox(value=getattr(llm_settings, "litellm_text2gql_tokens"), label="max_token"), ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - gr.Textbox(value=lambda: getattr(llm_settings, "litellm_text2gql_tokens"), label="max_token"), ->>>>>>> f42fa9b (feat(llm): use lambda) ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_text2gql_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_text2gql_op, inputs=llm_config_input) with gr.Accordion("3. Set up the Embedding.", open=False): embedding_dropdown = gr.Dropdown( -<<<<<<< HEAD choices=["openai", "litellm", "ollama/local"], -======= - choices=["openai", "litellm", "qianfan_wenxin", "ollama/local"], ->>>>>>> 38dce0b (feat(llm): vector db finished) value=llm_settings.embedding_type, label="Embedding", ) @@ -948,123 +567,67 @@ def embedding_settings(embedding_type): if embedding_type == "openai": with gr.Row(): embedding_config_input = [ -<<<<<<< HEAD -<<<<<<< HEAD gr.Textbox( - value=llm_settings.openai_embedding_api_key, + value=lambda: llm_settings.openai_embedding_api_key, label="api_key", type="password", ), gr.Textbox( - value=llm_settings.openai_embedding_api_base, + value=lambda: llm_settings.openai_embedding_api_base, label="api_base", ), gr.Textbox( - value=llm_settings.openai_embedding_model, + value=lambda: llm_settings.openai_embedding_model, label="model_name", ), -======= - gr.Textbox(value=llm_settings.openai_embedding_api_key, label="api_key", type="password"), - gr.Textbox(value=llm_settings.openai_embedding_api_base, label="api_base"), - gr.Textbox(value=llm_settings.openai_embedding_model, label="model_name"), - gr.Textbox(value=str(llm_settings.openai_embedding_model_dim), label="model_dim"), ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= gr.Textbox( - value=lambda: llm_settings.openai_embedding_api_key, label="api_key", type="password" + value=lambda: str(llm_settings.openai_embedding_model_dim), + label="model_dim", ), - gr.Textbox(value=lambda: llm_settings.openai_embedding_api_base, label="api_base"), - gr.Textbox(value=lambda: llm_settings.openai_embedding_model, label="model_name"), - gr.Textbox(value=lambda: str(llm_settings.openai_embedding_model_dim), label="model_dim"), ->>>>>>> f42fa9b (feat(llm): use lambda) ] elif embedding_type == "ollama/local": with gr.Row(): embedding_config_input = [ -<<<<<<< HEAD - gr.Textbox(value=llm_settings.ollama_embedding_host, label="host"), - gr.Textbox(value=str(llm_settings.ollama_embedding_port), label="port"), -<<<<<<< HEAD gr.Textbox( - value=llm_settings.ollama_embedding_model, - label="model_name", + value=lambda: llm_settings.ollama_embedding_host, + label="host", + ), + gr.Textbox( + value=lambda: str(llm_settings.ollama_embedding_port), + label="port", ), -======= - gr.Textbox(value=llm_settings.ollama_embedding_model, label="model_name"), - gr.Textbox(value=str(llm_settings.ollama_embedding_model_dim), label="model_dim"), -======= - gr.Textbox(value=lambda: llm_settings.ollama_embedding_host, label="host"), - gr.Textbox(value=lambda: str(llm_settings.ollama_embedding_port), label="port"), - gr.Textbox(value=lambda: llm_settings.ollama_embedding_model, label="model_name"), - gr.Textbox(value=lambda: str(llm_settings.ollama_embedding_model_dim), label="model_dim"), ->>>>>>> f42fa9b (feat(llm): use lambda) - ] - elif embedding_type == "qianfan_wenxin": - with gr.Row(): - embedding_config_input = [ gr.Textbox( - value=lambda: llm_settings.qianfan_embedding_api_key, label="api_key", type="password" + value=lambda: llm_settings.ollama_embedding_model, + label="model_name", ), -<<<<<<< HEAD - gr.Textbox(value=llm_settings.qianfan_embedding_model, label="model_name"), - gr.Textbox(value=str(llm_settings.qianfan_embedding_model_dim), label="model_dim"), ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= gr.Textbox( - value=lambda: llm_settings.qianfan_embedding_secret_key, label="secret_key", type="password" + value=lambda: str(llm_settings.ollama_embedding_model_dim), + label="model_dim", ), - gr.Textbox(value=lambda: llm_settings.qianfan_embedding_model, label="model_name"), - gr.Textbox(value=lambda: str(llm_settings.qianfan_embedding_model_dim), label="model_dim"), ->>>>>>> f42fa9b (feat(llm): use lambda) ] elif embedding_type == "litellm": with gr.Row(): embedding_config_input = [ gr.Textbox( -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - value=getattr(llm_settings, "litellm_embedding_api_key"), - label="api_key", - type="password", -======= - value=getattr(llm_settings, "litellm_embedding_api_key"), label="api_key", type="password" ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= value=lambda: getattr(llm_settings, "litellm_embedding_api_key"), -======= - value=lambda: getattr( - llm_settings, "litellm_embedding_api_key" - ), ->>>>>>> 3aeef7d (fix) label="api_key", type="password", ->>>>>>> f42fa9b (feat(llm): use lambda) ), gr.Textbox( - value=lambda: getattr( - llm_settings, "litellm_embedding_api_base" - ), + value=lambda: getattr(llm_settings, "litellm_embedding_api_base"), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=lambda: getattr( - llm_settings, "litellm_embedding_model" - ), + value=lambda: getattr(llm_settings, "litellm_embedding_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/embedding/supported_embedding", ), -<<<<<<< HEAD -======= gr.Textbox( - value=lambda: getattr( - llm_settings, "litellm_embedding_model_dim" - ), + value=lambda: getattr(llm_settings, "litellm_embedding_model_dim"), label="model_dim", type="text", ), ->>>>>>> 38dce0b (feat(llm): vector db finished) ] else: embedding_config_input = [ @@ -1091,47 +654,26 @@ def embedding_settings(embedding_type): @gr.render(inputs=[reranker_dropdown]) def reranker_settings(reranker_type): - llm_settings.reranker_type = ( - reranker_type if reranker_type != "None" else None - ) + llm_settings.reranker_type = reranker_type if reranker_type != "None" else None if reranker_type == "cohere": with gr.Row(): reranker_config_input = [ -<<<<<<< HEAD gr.Textbox( - value=llm_settings.reranker_api_key, + value=lambda: llm_settings.reranker_api_key, label="api_key", type="password", ), -<<<<<<< HEAD - gr.Textbox(value=llm_settings.reranker_model, label="model"), - gr.Textbox(value=llm_settings.cohere_base_url, label="base_url"), -======= - gr.Textbox(value=lambda: llm_settings.reranker_api_key, label="api_key", type="password"), gr.Textbox(value=lambda: llm_settings.reranker_model, label="model"), gr.Textbox(value=lambda: llm_settings.cohere_base_url, label="base_url"), ->>>>>>> f42fa9b (feat(llm): use lambda) -======= - gr.Textbox( - value=lambda: llm_settings.reranker_model, label="model" - ), - gr.Textbox( - value=lambda: llm_settings.cohere_base_url, label="base_url" - ), ->>>>>>> 3aeef7d (fix) ] elif reranker_type == "siliconflow": with gr.Row(): reranker_config_input = [ -<<<<<<< HEAD gr.Textbox( - value=llm_settings.reranker_api_key, + value=lambda: llm_settings.reranker_api_key, label="api_key", type="password", ), -======= - gr.Textbox(value=lambda: llm_settings.reranker_api_key, label="api_key", type="password"), ->>>>>>> f42fa9b (feat(llm): use lambda) gr.Textbox( value="BAAI/bge-reranker-v2-m3", label="model", @@ -1154,13 +696,7 @@ def reranker_settings(reranker_type): inputs=reranker_config_input, # pylint: disable=no-member ) -<<<<<<< HEAD -<<<<<<< HEAD -======= - with gr.Accordion("5. Set up the vector database.", open=False): -======= with gr.Accordion("5. Set up the vector engine.", open=False): ->>>>>>> a255aed (fix cycle import & add docs) engine_selector = gr.Dropdown( choices=["Faiss", "Milvus", "Qdrant"], value=index_settings.cur_vector_index, @@ -1170,7 +706,7 @@ def reranker_settings(reranker_type): fn=lambda engine: setattr(index_settings, "cur_vector_index", engine), inputs=[engine_selector], ) ->>>>>>> 38dce0b (feat(llm): vector db finished) + # The reason for returning this partial value is the functional need to refresh the ui return graph_config_input diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py index 3f8089b6d..82650b907 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py @@ -34,13 +34,9 @@ def create_other_block(): inp = gr.Textbox( value="g.V().limit(10)", label="Gremlin query", show_copy_button=True, lines=8 ) - out = gr.Code( - label="Output", language="json", elem_classes="code-container-show" - ) + out = gr.Code(label="Output", language="json", elem_classes="code-container-show") btn = gr.Button("Run Gremlin query") - btn.click( - fn=run_gremlin_query, inputs=[inp], outputs=out - ) # pylint: disable=no-member + btn.click(fn=run_gremlin_query, inputs=[inp], outputs=out) # pylint: disable=no-member gr.Markdown("---") with gr.Row(): @@ -55,9 +51,7 @@ def create_other_block(): inp = [] out = gr.Textbox(label="Init Graph Demo Result", show_copy_button=True) btn = gr.Button("(BETA) Init HugeGraph test data (🚧)") - btn.click( - fn=init_hg_test_data, inputs=inp, outputs=out - ) # pylint: disable=no-member + btn.click(fn=init_hg_test_data, inputs=inp, outputs=out) # pylint: disable=no-member @asynccontextmanager diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index de4e0b82c..2a817ada4 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -21,16 +21,17 @@ from typing import Any, AsyncGenerator, Literal, Optional, Tuple import gradio as gr -from hugegraph_llm.flows.scheduler import SchedulerSingleton import pandas as pd -<<<<<<< HEAD -from hugegraph_llm.config import resource_path, prompt, llm_settings -======= -from hugegraph_llm.config import huge_settings, index_settings, llm_settings, prompt, resource_path +from hugegraph_llm.config import ( + huge_settings, + index_settings, + llm_settings, + prompt, + resource_path, +) from hugegraph_llm.operators.graph_rag_task import RAGPipeline from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize ->>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.utils.decorators import with_task_id from hugegraph_llm.utils.log import log @@ -76,9 +77,6 @@ def rag_answer( gr.Warning("Please select at least one generate mode.") return "", "", "", "" -<<<<<<< HEAD - scheduler = SchedulerSingleton.get_instance() -======= rag = RAGPipeline() if vector_search: rag.query_vector_index(vector_index_str=index_settings.cur_vector_index) @@ -99,57 +97,29 @@ def rag_answer( near_neighbor_first=near_neighbor_first, topk_return_results=topk_return_results, ) - rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) + rag.synthesize_answer( + raw_answer, + vector_only_answer, + graph_only_answer, + graph_vector_answer, + answer_prompt, + ) ->>>>>>> 38dce0b (feat(llm): vector db finished) try: - # Select workflow by mode to avoid fetching the wrong pipeline from the pool - if graph_vector_answer or (graph_only_answer and vector_only_answer): - flow_key = "rag_graph_vector" - elif vector_only_answer: - flow_key = "rag_vector_only" - elif graph_only_answer: - flow_key = "rag_graph_only" - elif raw_answer: - flow_key = "rag_raw" - else: - raise RuntimeError("Unsupported flow type") - - res = scheduler.schedule_flow( - flow_key, + context = rag.run( + verbose=True, query=text, vector_search=vector_search, graph_search=graph_search, - raw_answer=raw_answer, - vector_only_answer=vector_only_answer, - graph_only_answer=graph_only_answer, - graph_vector_answer=graph_vector_answer, - graph_ratio=graph_ratio, - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - custom_related_information=custom_related_information, - answer_prompt=answer_prompt, - keywords_extract_prompt=keywords_extract_prompt, - gremlin_tmpl_num=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, max_graph_items=max_graph_items, - topk_return_results=topk_return_results, - vector_dis_threshold=vector_dis_threshold, - topk_per_keyword=topk_per_keyword, ) -<<<<<<< HEAD - if res.get("switch_to_bleu"): -======= if context.get("switch_to_bleu"): ->>>>>>> 3aeef7d (fix) - gr.Warning( - "Online reranker fails, automatically switches to local bleu rerank." - ) + gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") return ( - res.get("raw_answer", ""), - res.get("vector_only_answer", ""), - res.get("graph_only_answer", ""), - res.get("graph_vector_answer", ""), + context.get("raw_answer", ""), + context.get("vector_only_answer", ""), + context.get("graph_only_answer", ""), + context.get("graph_vector_answer", ""), ) except ValueError as e: log.critical(e) @@ -227,8 +197,6 @@ async def rag_answer_streaming( yield "", "", "", "" return -<<<<<<< HEAD -======= rag = RAGPipeline() if vector_search: rag.query_vector_index(vector_index_str=index_settings.cur_vector_index) @@ -246,64 +214,30 @@ async def rag_answer_streaming( ) # rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) ->>>>>>> 38dce0b (feat(llm): vector db finished) try: - # Select the specific streaming workflow - scheduler = SchedulerSingleton.get_instance() - if graph_vector_answer or (graph_only_answer and vector_only_answer): - flow_key = "rag_graph_vector" - elif vector_only_answer: - flow_key = "rag_vector_only" - elif graph_only_answer: - flow_key = "rag_graph_only" - elif raw_answer: - flow_key = "rag_raw" - else: - raise RuntimeError("Unsupported flow type") - - async for res in scheduler.schedule_stream_flow( - flow_key, + context = rag.run( + verbose=True, query=text, vector_search=vector_search, graph_search=graph_search, -<<<<<<< HEAD -======= ) if context.get("switch_to_bleu"): - gr.Warning( - "Online reranker fails, automatically switches to local bleu rerank." - ) + gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") answer_synthesize = AnswerSynthesize( ->>>>>>> 3aeef7d (fix) raw_answer=raw_answer, vector_only_answer=vector_only_answer, graph_only_answer=graph_only_answer, graph_vector_answer=graph_vector_answer, -<<<<<<< HEAD - graph_ratio=graph_ratio, - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - custom_related_information=custom_related_information, - answer_prompt=answer_prompt, - keywords_extract_prompt=keywords_extract_prompt, - gremlin_tmpl_num=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - ): - if res.get("switch_to_bleu"): -======= prompt_template=answer_prompt, ) async for context in answer_synthesize.run_streaming(context): if context.get("switch_to_bleu"): ->>>>>>> 3aeef7d (fix) - gr.Warning( - "Online reranker fails, automatically switches to local bleu rerank." - ) + gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") yield ( - res.get("raw_answer", ""), - res.get("vector_only_answer", ""), - res.get("graph_only_answer", ""), - res.get("graph_vector_answer", ""), + context.get("raw_answer", ""), + context.get("vector_only_answer", ""), + context.get("graph_only_answer", ""), + context.get("graph_vector_answer", ""), ) except ValueError as e: log.critical(e) @@ -368,23 +302,10 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): -<<<<<<< HEAD -<<<<<<< HEAD -======= ->>>>>>> 3aeef7d (fix) - raw_radio = gr.Radio( - choices=[True, False], value=False, label="Basic LLM Answer" - ) + raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") vector_only_radio = gr.Radio( choices=[True, False], value=False, label="Vector-only Answer" ) -<<<<<<< HEAD -======= - raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") - vector_only_radio = gr.Radio(choices=[True, False], value=False, label="Vector-only Answer") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= ->>>>>>> 3aeef7d (fix) with gr.Row(): graph_only_radio = gr.Radio( choices=[True, False], value=True, label="Graph-only Answer" @@ -452,7 +373,7 @@ def toggle_slider(enable): """## 2. (Batch) Back-testing ) > 1. Download the template file & fill in the questions you want to test. > 2. Upload the file & click the button to generate answers. (Preview shows the first 40 lines) - > 3. The answer options are the same as the above RAG/Q&A frame + > 3. The answer options are the same as the above RAG/Q&A frame """ ) tests_df_headers = [ @@ -466,9 +387,7 @@ def toggle_slider(enable): # FIXME: "demo" might conflict with the graph name, it should be modified. answers_path = os.path.join(resource_path, "demo", "questions_answers.xlsx") questions_path = os.path.join(resource_path, "demo", "questions.xlsx") - questions_template_path = os.path.join( - resource_path, "demo", "questions_template.xlsx" - ) + questions_template_path = os.path.join(resource_path, "demo", "questions_template.xlsx") def read_file_to_excel(file: Any, line_count: Optional[int] = None): df = None @@ -548,18 +467,12 @@ def several_rag_answer( file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" ) with gr.Column(): - test_template_file = os.path.join( - resource_path, "demo", "questions_template.xlsx" - ) + test_template_file = os.path.join(resource_path, "demo", "questions_template.xlsx") gr.File(value=test_template_file, label="Download Template File") - answer_max_line_count = gr.Number( - 1, label="Max Lines To Show", minimum=1, maximum=40 - ) + answer_max_line_count = gr.Number(1, label="Max Lines To Show", minimum=1, maximum=40) answers_btn = gr.Button("Generate Answer (Batch)", variant="primary") # TODO: Set individual progress bars for dataframe - qa_dataframe = gr.DataFrame( - label="Questions & Answers (Preview)", headers=tests_df_headers - ) + qa_dataframe = gr.DataFrame(label="Questions & Answers (Preview)", headers=tests_df_headers) answers_btn.click( several_rag_answer, inputs=[ @@ -577,12 +490,8 @@ def several_rag_answer( ], outputs=[qa_dataframe, gr.File(label="Download Answered File", min_width=40)], ) - questions_file.change( - read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count] - ) - answer_max_line_count.change( - change_showing_excel, answer_max_line_count, qa_dataframe - ) + questions_file.change(read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count]) + answer_max_line_count.change(change_showing_excel, answer_max_line_count, qa_dataframe) return ( inp, answer_prompt_input, diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 1d4407519..8fbb01c25 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -17,13 +17,9 @@ import json import os -<<<<<<< HEAD -from datetime import datetime from dataclasses import dataclass -from typing import Any, Tuple, Dict, Literal, Optional, List -======= -from typing import Any, Dict, Literal, Tuple, Union ->>>>>>> 38dce0b (feat(llm): vector db finished) +from datetime import datetime +from typing import Any, Dict, List, Literal, Optional, Tuple import gradio as gr import pandas as pd @@ -37,8 +33,7 @@ from hugegraph_llm.utils.embedding_utils import get_index_folder_name from hugegraph_llm.utils.hugegraph_utils import run_gremlin_query from hugegraph_llm.utils.log import log -<<<<<<< HEAD -from hugegraph_llm.flows.scheduler import SchedulerSingleton +from hugegraph_llm.utils.vector_index_utils import get_vector_index_class @dataclass @@ -60,16 +55,12 @@ def error(cls, message: str) -> "GremlinResult": @classmethod def success_result( -<<<<<<< HEAD cls, match_result: str, template_gremlin: str, raw_gremlin: str, template_exec: str, raw_exec: str, -======= - cls, match_result: str, template_gremlin: str, raw_gremlin: str, template_exec: str, raw_exec: str ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ) -> "GremlinResult": """Create a successful result""" return cls( @@ -80,9 +71,6 @@ def success_result( template_exec_result=template_exec, raw_exec_result=raw_exec, ) -======= -from hugegraph_llm.utils.vector_index_utils import get_vector_index_class ->>>>>>> 38dce0b (feat(llm): vector db finished) def store_schema(schema, question, gremlin_prompt): @@ -98,23 +86,8 @@ def store_schema(schema, question, gremlin_prompt): def build_example_vector_index(temp_file) -> dict: -<<<<<<< HEAD -<<<<<<< HEAD - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - index_path = os.path.join(resource_path, folder_name, "gremlin_examples") - if not os.path.exists(index_path): - os.makedirs(index_path) -======= - vector_index = get_vector_index_class(index_settings.now_vector_index) -<<<<<<< HEAD - assert vector_index, 'vector db name is error' ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= -======= vector_index = get_vector_index_class(index_settings.cur_vector_index) ->>>>>>> a255aed (fix cycle import & add docs) assert vector_index, "vector db name is error" ->>>>>>> 87ee5d3 (style: format code with black line-length 120) if temp_file is None: full_path = os.path.join(resource_path, "demo", "text2gremlin.csv") else: @@ -123,23 +96,12 @@ def build_example_vector_index(temp_file) -> dict: timestamp = datetime.now().strftime("%Y%m%d%H%M%S") _, file_name = os.path.split(f"{name}_{timestamp}{ext}") log.info("Copying file to: %s", file_name) -<<<<<<< HEAD + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) target_file = os.path.join(resource_path, folder_name, "gremlin_examples", file_name) -======= - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - target_file = os.path.join( - resource_path, folder_name, "gremlin_examples", file_name - ) ->>>>>>> 3aeef7d (fix) try: import shutil -<<<<<<< HEAD -======= os.makedirs(os.path.dirname(target_file), exist_ok=True) ->>>>>>> 87ee5d3 (style: format code with black line-length 120) shutil.copy2(full_path, target_file) log.info("Successfully copied file to: %s", target_file) except (OSError, IOError) as e: @@ -161,17 +123,8 @@ def build_example_vector_index(temp_file) -> dict: return builder.example_index_build(examples, vector_index=vector_index).run() -<<<<<<< HEAD def _process_schema(schema, generator, sm): """Process and validate schema input""" -======= -def gremlin_generate( - inp, example_num, schema, gremlin_prompt -) -> Union[tuple[str, str], tuple[str, Any, Any, Any, Any]]: - vector_index = get_vector_index_class(index_settings.now_vector_index) - generator = GremlinGenerator(llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding()) - sm = SchemaManager(graph_name=schema) ->>>>>>> 38dce0b (feat(llm): vector db finished) short_schema = False if not schema: return None, short_schema @@ -230,37 +183,23 @@ def _execute_queries(context, output_types): def gremlin_generate( - inp, example_num, schema, gremlin_prompt, requested_outputs: Optional[List[str]] = None + inp, + example_num, + schema, + gremlin_prompt, + requested_outputs: Optional[List[str]] = None, ) -> GremlinResult: -<<<<<<< HEAD -<<<<<<< HEAD - generator = GremlinGenerator( - llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() - ) -======= - vector_index = get_vector_index_class(index_settings.now_vector_index) -======= vector_index = get_vector_index_class(index_settings.cur_vector_index) -<<<<<<< HEAD ->>>>>>> a255aed (fix cycle import & add docs) - generator = GremlinGenerator(llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding()) ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= generator = GremlinGenerator( llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() ) ->>>>>>> 3aeef7d (fix) sm = SchemaManager(graph_name=schema) processed_schema, short_schema = _process_schema(schema, generator, sm) if processed_schema is None and short_schema is None: - return GremlinResult.error( - "Invalid JSON schema, please check the format carefully." - ) + return GremlinResult.error("Invalid JSON schema, please check the format carefully.") - updated_schema = ( - sm.simple_schema(processed_schema) if short_schema else processed_schema - ) + updated_schema = sm.simple_schema(processed_schema) if short_schema else processed_schema store_schema(str(updated_schema), inp, gremlin_prompt) output_types = _configure_output_types(requested_outputs) @@ -292,11 +231,7 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = { - key: vertex[key] - for key in ["id", "name", "properties"] - if key in vertex - } + new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) @@ -315,48 +250,17 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: def gremlin_generate_for_ui(inp, example_num, schema, gremlin_prompt): """UI wrapper for gremlin_generate that returns tuple for Gradio compatibility""" - # Execute via scheduler - try: - res = SchedulerSingleton.get_instance().schedule_flow( - "text2gremlin", - inp, - int(example_num) if isinstance(example_num, (int, float, str)) else 2, - schema, - gremlin_prompt, - [ - "match_result", - "template_gremlin", - "raw_gremlin", - "template_execution_result", - "raw_execution_result", - ], - ) - except Exception as e: # pylint: disable=broad-except - log.error("UI text2gremlin error: %s", e) - return json.dumps({"error": str(e)}, ensure_ascii=False), "", "", "", "" - - # Backward-compatible mapping for outputs - match_result = res.get("match_result", []) - match_result_str = ( - json.dumps(match_result, ensure_ascii=False, indent=2) - if isinstance(match_result, (list, dict)) - else str(match_result) - ) + result = gremlin_generate(inp, example_num, schema, gremlin_prompt) + + if not result.success: + return result.match_result, "", "", "", "" return ( -<<<<<<< HEAD - match_result_str, - res.get("template_gremlin", "") or "", - res.get("raw_gremlin", "") or "", - res.get("template_execution_result", "") or "", - res.get("raw_execution_result", "") or "", -======= result.match_result, result.template_gremlin or "", result.raw_gremlin or "", result.template_exec_result or "", result.raw_exec_result or "", ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ) @@ -378,27 +282,23 @@ def create_text2gremlin_block() -> Tuple: with gr.Row(): btn = gr.Button("Build Example Vector Index", variant="primary") - btn.click( - build_example_vector_index, inputs=[file], outputs=[out] - ) # pylint: disable=no-member + btn.click(build_example_vector_index, inputs=[file], outputs=[out]) # pylint: disable=no-member gr.Markdown("## Nature Language To Gremlin") with gr.Row(): with gr.Column(scale=1): input_box = gr.Textbox( - value=prompt.default_question, label="Nature Language Query", show_copy_button=True + value=prompt.default_question, + label="Nature Language Query", + show_copy_button=True, ) match = gr.Code( label="Similar Template (TopN)", language="javascript", elem_classes="code-container-show", ) - initialized_out = gr.Textbox( - label="Gremlin With Template", show_copy_button=True - ) - raw_out = gr.Textbox( - label="Gremlin Without Template", show_copy_button=True - ) + initialized_out = gr.Textbox(label="Gremlin With Template", show_copy_button=True) + raw_out = gr.Textbox(label="Gremlin Without Template", show_copy_button=True) tmpl_exec_out = gr.Code( label="Query With Template Output", language="json", @@ -415,7 +315,10 @@ def create_text2gremlin_block() -> Tuple: minimum=0, maximum=10, step=1, value=2, label="Number of refer examples" ) schema_box = gr.Textbox( - value=prompt.text2gql_graph_schema, label="Schema", lines=2, show_copy_button=True + value=prompt.text2gql_graph_schema, + label="Schema", + lines=2, + show_copy_button=True, ) prompt_box = gr.Textbox( value=prompt.gremlin_generate_prompt, @@ -449,18 +352,7 @@ def graph_rag_recall( store_schema(prompt.text2gql_graph_schema, query, gremlin_prompt) rag = RAGPipeline() rag.extract_keywords().keywords_to_vid( -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -======= - vector_index=index_settings.now_vector_index, ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= - vector_index_str=index_settings.now_vector_index, ->>>>>>> dd3b085 (feat(llm): nexpected-keyword-arg,unused-import) -======= vector_index_str=index_settings.cur_vector_index, ->>>>>>> a255aed (fix cycle import & add docs) vector_dis_threshold=vector_dis_threshold, topk_per_keyword=topk_per_keyword, ) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 92ab0b895..f011e2082 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -90,9 +90,7 @@ def generate_prompt_for_ui(source_text, scenario, example_name): def load_example_names(): """Load all candidate examples""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "prompt_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return [example.get("name", "Unnamed example") for example in examples] @@ -127,22 +125,16 @@ def load_query_examples(): >>>>>>> 3aeef7d (fix) ) if language.upper() == "CN": - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples_CN.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples_CN.json") else: - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) except (FileNotFoundError, json.JSONDecodeError): try: - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -153,9 +145,7 @@ def load_query_examples(): def load_schema_fewshot_examples(): """Load few-shot examples from a JSON file""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "schema_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "schema_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -577,9 +567,7 @@ def create_vector_graph_block(): max_lines=29, ) - out = gr.Code( - label="Output Info", language="json", elem_classes="code-container-edit" - ) + out = gr.Code(label="Output Info", language="json", elem_classes="code-container-edit") with gr.Row(): with gr.Accordion("Get RAG Info", open=False): @@ -628,9 +616,7 @@ def create_vector_graph_block(): store_prompt, inputs=[input_text, input_schema, info_extract_template], ) - vector_import_bt.click( - build_vector_index, inputs=[input_file, input_text], outputs=out - ).then( + vector_import_bt.click(build_vector_index, inputs=[input_file, input_text], outputs=out).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) @@ -658,17 +644,15 @@ def create_vector_graph_block(): inputs=[input_text, input_schema, info_extract_template], ) - graph_loading_bt.click( - import_graph_data, inputs=[out, input_schema], outputs=[out] - ).then(update_vid_embedding).then( + graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( + update_vid_embedding + ).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) build_schema_bt.click( - lambda it, qe, fs: extract_graph( - [], it, prompt.graph_schema, prompt.extract_graph_prompt - ), + lambda it, qe, fs: extract_graph([], it, prompt.graph_schema, prompt.extract_graph_prompt), inputs=[input_text, query_example, few_shot], outputs=[input_schema], ).then( diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py index fae5ca860..b168b8c39 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -73,15 +73,11 @@ def __init__( def _create_collection(self): """Create a new collection in Milvus.""" - id_field = FieldSchema( - name="id", dtype=DataType.INT64, is_primary=True, auto_id=True - ) + id_field = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True) vector_field = FieldSchema( name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embed_dim ) - property_field = FieldSchema( - name="property", dtype=DataType.VARCHAR, max_length=65535 - ) + property_field = FieldSchema(name="property", dtype=DataType.VARCHAR, max_length=65535) original_id_field = FieldSchema(name="original_id", dtype=DataType.INT64) schema = CollectionSchema( diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index 14b97fbff..ca7761ecd 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -28,9 +28,7 @@ class QdrantVectorIndex(VectorStoreBase): - def __init__( - self, name: str, host: str, port: int, api_key=None, embed_dim: int = 1024 - ): + def __init__(self, name: str, host: str, port: int, api_key=None, embed_dim: int = 1024): self.embed_dim = embed_dim self.host = host self.port = port @@ -117,9 +115,7 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: return remove_num - def search( - self, query_vector: List[float], top_k: int = 5, dis_threshold: float = 0.9 - ): + def search(self, query_vector: List[float], top_k: int = 5, dis_threshold: float = 0.9): search_result = self.client.search( collection_name=self.name, query_vector=query_vector, limit=top_k ) diff --git a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py index 5d98ebdf8..f13e11e6f 100644 --- a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py +++ b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py @@ -41,9 +41,7 @@ async def dispatch(self, request: Request, call_next): unit = "s" response.headers["X-Process-Time"] = f"{process_time:.2f} {unit}" - log.info( - "Request process time: %.2f ms, code=%d", process_time, response.status_code - ) + log.info("Request process time: %.2f ms, code=%d", process_time, response.status_code) log.info( "%s - Args: %s, IP: %s, URL: %s", request.method, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index 8e6af2774..d96840911 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -65,9 +65,7 @@ def __init__(self): def get_embedding(self): if self.embedding_type == "openai": - assert ( - llm_settings.openai_embedding_model_dim - ), "openai_embedding_model_dim is need" + assert llm_settings.openai_embedding_model_dim, "openai_embedding_model_dim is need" return OpenAIEmbedding( embedding_dimension=llm_settings.openai_embedding_model_dim, model_name=llm_settings.openai_embedding_model, @@ -75,9 +73,7 @@ def get_embedding(self): api_base=llm_settings.openai_embedding_api_base, ) if self.embedding_type == "ollama/local": - assert ( - llm_settings.ollama_embedding_model_dim - ), "ollama_embedding_model_dim is need" + assert llm_settings.ollama_embedding_model_dim, "ollama_embedding_model_dim is need" return OllamaEmbedding( <<<<<<< HEAD model_name=llm_settings.ollama_embedding_model, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index 15d928286..6e30cca71 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -91,7 +91,5 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] A list of embedding vectors, where each vector is a list of floats. The order of embeddings should match the order of input texts. """ - response = await self.aclient.embeddings.create( - input=texts, model=self.model_name - ) + response = await self.aclient.embeddings.create(input=texts, model=self.model_name) return [data.embedding for data in response.data] 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 7e1eaab68..9121fca09 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py @@ -173,8 +173,4 @@ def get_text2gql_llm(self): if __name__ == "__main__": client = LLMs().get_chat_llm() print(client.generate(prompt="What is the capital of China?")) - print( - client.generate( - messages=[{"role": "user", "content": "What is the capital of China?"}] - ) - ) + print(client.generate(messages=[{"role": "user", "content": "What is the capital of China?"}])) diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py index c15c5440e..199384b12 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py @@ -29,9 +29,7 @@ class OllamaClient(BaseLLM): """LLM wrapper should take in a prompt and return a string.""" - def __init__( - self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwargs - ): + def __init__(self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwargs): self.model = model self.client = ollama.Client(host=f"http://{host}:{port}", **kwargs) self.async_client = ollama.AsyncClient(host=f"http://{host}:{port}", **kwargs) @@ -101,9 +99,7 @@ def generate_streaming( for chunk in self.client.chat(model=self.model, messages=messages, stream=True): if not chunk["message"]: - log.debug( - "Received empty chunk['message'] in streaming chunk: %s", chunk - ) + log.debug("Received empty chunk['message'] in streaming chunk: %s", chunk) continue token = chunk["message"]["content"] if on_token_callback: diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py index 52d624941..e1088c890 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py @@ -52,9 +52,7 @@ def __init__( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type( - (RateLimitError, APIConnectionError, APITimeoutError) - ), + retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), ) def generate( self, @@ -89,9 +87,7 @@ def generate( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type( - (RateLimitError, APIConnectionError, APITimeoutError) - ), + retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), ) async def agenerate( self, @@ -126,9 +122,7 @@ async def agenerate( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type( - (RateLimitError, APIConnectionError, APITimeoutError) - ), + retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), ) def generate_streaming( self, diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py index 9886aa0ca..3bf481ce2 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py @@ -57,9 +57,7 @@ def get_rerank_lists( "top_n": top_n, "documents": documents, } - response = requests.post( - url, headers=headers, json=payload, timeout=(1.0, 10.0) - ) + response = requests.post(url, headers=headers, json=payload, timeout=(1.0, 10.0)) response.raise_for_status() # Raise an error for bad status codes results = response.json()["results"] sorted_docs = [documents[item["index"]] for item in results] diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py index da8a9f7b7..e4a9b550a 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py @@ -58,9 +58,7 @@ def get_rerank_lists( "content-type": Constants.HEADER_CONTENT_TYPE, "authorization": f"Bearer {self.api_key}", } - response = requests.post( - url, json=payload, headers=headers, timeout=(1.0, 10.0) - ) + response = requests.post(url, json=payload, headers=headers, timeout=(1.0, 10.0)) response.raise_for_status() # Raise an error for bad status codes results = response.json()["results"] sorted_docs = [documents[item["index"]] for item in results] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index bd2479817..47b0f060f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -59,12 +59,8 @@ def _validate_schema(self, schema: Dict[str, Any]) -> None: check_type(schema, dict, "Input data is not a dictionary.") if "vertexlabels" not in schema or "edgelabels" not in schema: log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") - check_type( - schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." - ) - check_type( - schema["edgelabels"], list, "'edgelabels' in input data is not a list." - ) + check_type(schema["vertexlabels"], list, "'vertexlabels' in input data is not a list.") + check_type(schema["edgelabels"], list, "'edgelabels' in input data is not a list.") def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): property_labels = schema.get("propertykeys", []) @@ -82,19 +78,13 @@ def _process_vertex_labels( for vertex_label in schema["vertexlabels"]: self._validate_vertex_label(vertex_label) properties = vertex_label["properties"] - primary_keys = self._process_keys( - vertex_label, "primary_keys", properties[:1] - ) + primary_keys = self._process_keys(vertex_label, "primary_keys", properties[:1]) if len(primary_keys) == 0: log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") vertex_label["primary_keys"] = primary_keys - nullable_keys = self._process_keys( - vertex_label, "nullable_keys", properties[1:] - ) + nullable_keys = self._process_keys(vertex_label, "nullable_keys", properties[1:]) vertex_label["nullable_keys"] = nullable_keys - self._add_missing_properties( - properties, property_labels, property_label_set - ) + self._add_missing_properties(properties, property_labels, property_label_set) def _process_edge_labels( self, schema: Dict[str, Any], property_labels: list, property_label_set: set @@ -102,17 +92,13 @@ def _process_edge_labels( for edge_label in schema["edgelabels"]: self._validate_edge_label(edge_label) properties = edge_label.get("properties", []) - self._add_missing_properties( - properties, property_labels, property_label_set - ) + self._add_missing_properties(properties, property_labels, property_label_set) def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") if "name" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'name'.") - check_type( - vertex_label["name"], str, "'name' in vertex_label is not of correct type." - ) + check_type(vertex_label["name"], str, "'name' in vertex_label is not of correct type.") if "properties" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'properties'.") check_type( @@ -133,9 +119,7 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: log_and_raise( "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." ) - check_type( - edge_label["name"], str, "'name' in edge_label is not of correct type." - ) + check_type(edge_label["name"], str, "'name' in edge_label is not of correct type.") check_type( edge_label["source_label"], str, @@ -147,13 +131,9 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: "'target_label' in edge_label is not of correct type.", ) - def _process_keys( - self, label: Dict[str, Any], key_type: str, default_keys: list - ) -> list: + def _process_keys(self, label: Dict[str, Any], key_type: str, default_keys: list) -> list: keys = label.get(key_type, default_keys) - check_type( - keys, list, f"'{key_type}' in {label['name']} is not of correct type." - ) + check_type(keys, list, f"'{key_type}' in {label['name']} is not of correct type.") new_keys = [key for key in keys if key in label["properties"]] return new_keys diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py index a257ccc4c..dc5b15e00 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py @@ -138,8 +138,7 @@ def _rerank_with_vertex_degree( if self.method == "bleu": vertex_rerank_res = [ - _bleu_rerank(query, vertex_degree) + [""] - for vertex_degree in vertex_degree_list + _bleu_rerank(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list ] depth = len(vertex_degree_list) @@ -147,14 +146,11 @@ def _rerank_with_vertex_degree( if result not in knowledge_with_degree: knowledge_with_degree[result] = [result] + [""] * (depth - 1) if len(knowledge_with_degree[result]) < depth: - knowledge_with_degree[result] += [""] * ( - depth - len(knowledge_with_degree[result]) - ) + knowledge_with_degree[result] += [""] * (depth - len(knowledge_with_degree[result])) def sort_key(res: str) -> Tuple[int, ...]: return tuple( - vertex_rerank_res[i].index(knowledge_with_degree[res][i]) - for i in range(depth) + vertex_rerank_res[i].index(knowledge_with_degree[res][i]) for i in range(depth) ) sorted_results = sorted(results, key=sort_key) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py index c23bf7735..797ea70ae 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py @@ -71,9 +71,7 @@ def get_cache_dir() -> str: # Windows (hopefully) else: - local = os.environ.get("LOCALAPPDATA", None) or os.path.expanduser( - "~\\AppData\\Local" - ) + local = os.environ.get("LOCALAPPDATA", None) or os.path.expanduser("~\\AppData\\Local") path = Path(local, "hugegraph_llm") if not os.path.exists(path): diff --git a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py index 0d160e1e5..0d9967020 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py @@ -78,11 +78,7 @@ def _filter_keywords( sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: results.update( - { - w - for w in sub_tokens - if w not in NLTKHelper().stopwords(lang=self._language) - } + {w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language)} ) return list(results) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index fa6b79f91..6eb805271 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -42,9 +42,7 @@ class RAGPipeline: querying graph databases and vector indices, merging and re-ranking results, and generating answers. """ - def __init__( - self, llm: Optional[BaseLLM] = None, embedding: Optional[BaseEmbedding] = None - ): + def __init__(self, llm: Optional[BaseLLM] = None, embedding: Optional[BaseEmbedding] = None): """ Initialize the RAGPipeline with optional LLM and embedding models. diff --git a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py index 52f50fdd6..b205bb798 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py @@ -41,15 +41,11 @@ def clear(self): def example_index_build(self, examples, vector_index: type[VectorStoreBase]): self.operators.append( - BuildGremlinExampleIndex( - self.embedding, examples, vector_index=vector_index - ) + BuildGremlinExampleIndex(self.embedding, examples, vector_index=vector_index) ) return self - def import_schema( - self, from_hugegraph=None, from_extraction=None, from_user_defined=None - ): + def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): if from_hugegraph: self.operators.append(SchemaManager(from_hugegraph)) elif from_user_defined: @@ -61,17 +57,13 @@ def import_schema( return self def example_index_query(self, num_examples, vector_index: type[VectorStoreBase]): - self.operators.append( - GremlinExampleIndexQuery(vector_index, self.embedding, num_examples) - ) + self.operators.append(GremlinExampleIndexQuery(vector_index, self.embedding, num_examples)) return self def gremlin_generate_synthesize( self, schema, gremlin_prompt: Optional[str] = None, vertices: Optional[List[str]] = None ): - self.operators.append( - GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt) - ) + self.operators.append(GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt)) return self def print_result(self): 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 aa886c0af..52626b72b 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 @@ -50,9 +50,7 @@ def run(self, data: dict) -> Dict[str, Any]: if not schema: # TODO: ensure the function works correctly (update the logic later) self.schema_free_mode(data.get("triples", [])) - log.warning( - "Using schema_free mode, could try schema_define mode for better effect!" - ) + log.warning("Using schema_free mode, could try schema_define mode for better effect!") else: self.init_schema_if_need(schema) self.load_into_graph(vertices, edges, schema) @@ -68,9 +66,7 @@ def _set_default_property(self, key, input_properties, property_label_map): # list or set default_value = [] input_properties[key] = default_value - log.warning( - "Property '%s' missing in vertex, set to '%s' for now", key, default_value - ) + log.warning("Property '%s' missing in vertex, set to '%s' for now", key, default_value) def _handle_graph_creation(self, func, *args, **kwargs): try: @@ -82,17 +78,11 @@ def _handle_graph_creation(self, func, *args, **kwargs): log.error("Error on creating: %s, %s", args, e) return None - def load_into_graph( - self, vertices, edges, schema - ): # pylint: disable=too-many-statements + def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many-statements # pylint: disable=R0912 (too-many-branches) - vertex_label_map = { - v_label["name"]: v_label for v_label in schema["vertexlabels"] - } + 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"] - } + property_label_map = {p_label["name"]: p_label for p_label in schema["propertykeys"]} for vertex in vertices: input_label = vertex["label"] @@ -108,9 +98,7 @@ def load_into_graph( vertex_label = vertex_label_map[input_label] primary_keys = vertex_label["primary_keys"] nullable_keys = vertex_label.get("nullable_keys", []) - non_null_keys = [ - key for key in vertex_label["properties"] if key not in nullable_keys - ] + non_null_keys = [key for key in vertex_label["properties"] if key not in nullable_keys] has_problem = False # 2. Handle primary-keys mode vertex @@ -142,9 +130,7 @@ def load_into_graph( # 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 - ) + 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(): @@ -181,9 +167,7 @@ def load_into_graph( 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 - ) + self._handle_graph_creation(self.client.graph().addEdge, label, start, end, properties) def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -207,20 +191,18 @@ def init_schema_if_need(self, schema: dict): source_vertex_label = edge["source_label"] target_vertex_label = edge["target_label"] properties = edge["properties"] - self.schema.edgeLabel(edge_label).sourceLabel( - source_vertex_label - ).targetLabel(target_vertex_label).properties(*properties).nullableKeys( - *properties - ).ifNotExist().create() + self.schema.edgeLabel(edge_label).sourceLabel(source_vertex_label).targetLabel( + target_vertex_label + ).properties(*properties).nullableKeys(*properties).ifNotExist().create() def schema_free_mode(self, 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("name").ifNotExist().create() + self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel("vertex").properties( + "name" + ).ifNotExist().create() self.schema.indexLabel("vertexByName").onV("vertex").by( "name" @@ -280,9 +262,7 @@ def _set_property_data_type(self, property_key, data_type): log.warning("UUID type is not supported, use text instead") property_key.asText() else: - log.error( - "Unknown data type %s for property_key %s", data_type, property_key - ) + log.error("Unknown data type %s for property_key %s", data_type, property_key) def _set_property_cardinality(self, property_key, cardinality): if cardinality == PropertyCardinality.SINGLE: @@ -292,13 +272,9 @@ def _set_property_cardinality(self, property_key, cardinality): elif cardinality == PropertyCardinality.SET: property_key.valueSet() else: - log.error( - "Unknown cardinality %s for property_key %s", cardinality, property_key - ) + 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: + def _check_property_data_type(self, data_type: str, cardinality: str, value) -> bool: if cardinality in ( PropertyCardinality.LIST.value, PropertyCardinality.SET.value, @@ -328,9 +304,7 @@ def _check_single_data_type(self, data_type: str, value) -> bool: if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below - if ( - data_type == PropertyDataType.DATE.value - ): # the format should be "yyyy-MM-dd" + if data_type == PropertyDataType.DATE.value: # the format should be "yyyy-MM-dd" import re return isinstance(value, str) and re.match(r"^\d{4}-\d{2}-\d{2}$", value) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py index 73c9530df..4c4c167c4 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py @@ -47,7 +47,5 @@ def res = [:]; result = self.graph.gremlin().exec(groovy_code)["data"] if isinstance(result, list) and len(result) > 0: - graph_summary.update( - {key: result[i].get(key) for i, key in enumerate(keys)} - ) + graph_summary.update({key: result[i].get(key) for i, key in enumerate(keys)}) return graph_summary diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py index 52399d99e..d9542cd39 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py @@ -132,9 +132,7 @@ def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: query_embedding = context.get("query_embedding") self._gremlin_generator.clear() - self._gremlin_generator.example_index_query( - num_examples=self._num_gremlin_generate_example - ) + self._gremlin_generator.example_index_query(num_examples=self._num_gremlin_generate_example) gremlin_response = self._gremlin_generator.gremlin_generate_synthesize( context["simple_schema"], vertices=vertices, gremlin_prompt=self._gremlin_prompt ).run(query=query, query_embedding=query_embedding) @@ -148,9 +146,7 @@ def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: result = self._client.gremlin().exec(gremlin=gremlin)["data"] if result == [None]: result = [] - context["graph_result"] = [ - json.dumps(item, ensure_ascii=False) for item in result - ] + context["graph_result"] = [json.dumps(item, ensure_ascii=False) for item in result] if context["graph_result"]: context["graph_result_flag"] = 1 context["graph_context_head"] = ( @@ -228,9 +224,7 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: "Unable to find vid, downgraded to property query, please confirm if it meets expectation." ) - paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)[ - "data" - ] + paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)["data"] graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( self._format_graph_query_result(query_paths=paths) ) @@ -352,18 +346,14 @@ def _process_vertex( use_id_to_match: bool, v_cache: Set[str], ) -> Tuple[str, int, int]: - matched_str = ( - item["id"] if use_id_to_match else item["props"][self._prop_to_match] - ) + matched_str = item["id"] if use_id_to_match else item["props"][self._prop_to_match] if matched_str in node_cache: flat_rel = flat_rel[:-prior_edge_str_len] return flat_rel, prior_edge_str_len, depth node_cache.add(matched_str) props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'v')}" - for k, v in item["props"].items() - if v + f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v ) # TODO: we may remove label id or replace with label name @@ -388,9 +378,7 @@ def _process_edge( e_cache: Set[Tuple[str, str, str]], ) -> Tuple[str, int]: props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'e')}" - for k, v in item["props"].items() - if v + f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v ) props_str = f"{{{props_str}}}" if props_str else "" prev_matched_str = ( @@ -407,9 +395,7 @@ def _process_edge( edge_label = item["label"] edge_str = ( - f"--[{edge_label}]-->" - if item["outV"] == prev_matched_str - else f"<--[{edge_label}]--" + f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" ) path_str += edge_str prior_edge_str_len = len(edge_str) @@ -427,20 +413,14 @@ def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: schema = self._get_graph_schema() vertex_props_str, edge_props_str = schema.split("\n")[:2] # TODO: rename to vertex (also need update in the schema) - vertex_props_str = ( - vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") - ) - edge_props_str = ( - edge_props_str[len("Edge properties: ") :].strip("[").strip("]") - ) + vertex_props_str = vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") + edge_props_str = edge_props_str[len("Edge properties: ") :].strip("[").strip("]") vertex_labels = self._extract_label_names(vertex_props_str) edge_labels = self._extract_label_names(edge_props_str) return vertex_labels, edge_labels @staticmethod - def _extract_label_names( - source: str, head: str = "name: ", tail: str = ", " - ) -> List[str]: + def _extract_label_names(source: str, head: str = "name: ", tail: str = ", ") -> List[str]: result = [] for s in source.split(head): end = s.find(tail) @@ -466,9 +446,7 @@ def _get_graph_schema(self, refresh: bool = False) -> str: log.debug("Link(Relation): %s", relationships) return self._schema - def _limit_property_query( - self, value: Optional[str], item_type: str - ) -> Optional[str]: + def _limit_property_query(self, value: Optional[str], item_type: str) -> Optional[str]: # NOTE: we skip the filter for list/set type (e.g., list of string, add it if needed) if not self._limit_property or not isinstance(value, str): return value 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 e9bccc2f4..2f0643a77 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 @@ -40,9 +40,7 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: new_vertex = { - key: vertex[key] - for key in ["id", "name", "properties"] - if key in vertex + key: vertex[key] for key in ["id", "name", "properties"] if key in vertex } mini_schema["vertexlabels"].append(new_vertex) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index d41f99a37..4f288a9de 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -62,14 +62,10 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: # !: We have assumed that self.example is not empty queries = [example["query"] for example in self.examples] # TODO: refactor function chain async to avoid blocking - examples_embedding = asyncio.run( - get_embeddings_parallel(self.embedding, queries) - ) + examples_embedding = asyncio.run(get_embeddings_parallel(self.embedding, queries)) embed_dim = len(examples_embedding[0]) if len(self.examples) > 0: - vector_index = self.vector_index.from_name( - embed_dim, self.vector_index_name - ) + vector_index = self.vector_index.from_name(embed_dim, self.vector_index_name) vector_index.add(examples_embedding, self.examples) <<<<<<< HEAD vector_index.to_index_file(self.index_dir, self.filename_prefix) 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 ba2b1e379..4c46b8906 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 @@ -15,139 +15,68 @@ # specific language governing permissions and limitations # under the License. - import asyncio from typing import Any, Dict -<<<<<<< HEAD -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex -from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager -from hugegraph_llm.utils.embedding_utils import ( - get_embeddings_parallel, - get_filename_prefix, - get_index_folder_name, -) -======= from tqdm import tqdm from hugegraph_llm.config import huge_settings from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager ->>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.utils.log import log class BuildSemanticIndex: -<<<<<<< HEAD -<<<<<<< HEAD - def __init__(self, embedding: BaseEmbedding): - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(embedding, "model_name", None) - ) - self.vid_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) -======= - def __init__(self, embedding: BaseEmbedding, vector_index: type[VectorStoreBase]): -<<<<<<< HEAD - self.vid_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids") ->>>>>>> 38dce0b (feat(llm): vector db finished) -======= 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") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= self.vid_index = vector_index.from_name( embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids" ) ->>>>>>> 3aeef7d (fix) self.embedding = embedding self.sm = SchemaManager(huge_settings.graph_name) def _extract_names(self, vertices: list[str]) -> list[str]: return [v.split(":")[1] for v in vertices] -<<<<<<< HEAD -======= async def _get_embeddings_parallel(self, vids: list[str]) -> list[Any]: sem = asyncio.Semaphore(10) batch_size = 1000 async def get_embeddings_with_semaphore(vid_list: list[str]) -> Any: - # Executes sync embedding method in a thread pool via loop.run_in_executor, combining async programming - # with multi-threading capabilities. - # This pattern avoids blocking the event loop and prepares for a future fully async pipeline. async with sem: loop = asyncio.get_running_loop() return await loop.run_in_executor( None, self.embedding.get_texts_embeddings, vid_list ) -<<<<<<< HEAD - # Split vids into batches of size batch_size vid_batches = [vids[i : i + batch_size] for i in range(0, len(vids), batch_size)] - - # Create tasks for each batch -======= - vid_batches = [ - vids[i : i + batch_size] for i in range(0, len(vids), batch_size) - ] ->>>>>>> 3aeef7d (fix) tasks = [get_embeddings_with_semaphore(batch) for batch in vid_batches] embeddings = [] with tqdm(total=len(tasks)) as pbar: for future in asyncio.as_completed(tasks): batch_embeddings = await future - embeddings.extend(batch_embeddings) # Extend the list with batch results + embeddings.extend(batch_embeddings) pbar.update(1) return embeddings ->>>>>>> 38dce0b (feat(llm): vector db finished) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: vertexlabels = self.sm.schema.getSchema()["vertexlabels"] - all_pk_flag = all( - data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels - ) + all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) past_vids = self.vid_index.get_all_properties() # only support Faiss # TODO: We should build vid vector index separately, especially when the vertices may be very large -<<<<<<< HEAD -<<<<<<< HEAD - -======= ->>>>>>> 38dce0b (feat(llm): vector db finished) present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py -======= - present_vids = context[ - "vertices" - ] # Warning: data truncated by fetch_graph_data.py ->>>>>>> 3aeef7d (fix) removed_vids = set(past_vids) - set(present_vids) removed_num = self.vid_index.remove(removed_vids) added_vids = list(set(present_vids) - set(past_vids)) if added_vids: -<<<<<<< HEAD vids_to_process = self._extract_names(added_vids) if all_pk_flag else added_vids - added_embeddings = asyncio.run(get_embeddings_parallel(self.embedding, vids_to_process)) -======= - vids_to_process = ( - self._extract_names(added_vids) if all_pk_flag else added_vids - ) - added_embeddings = asyncio.run( - self._get_embeddings_parallel(vids_to_process) - ) ->>>>>>> 3aeef7d (fix) + 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) -<<<<<<< HEAD - self.vid_index.to_index_file(self.index_dir, self.filename_prefix) + self.vid_index.save_index_by_name(huge_settings.graph_name, "graph_vids") else: log.debug("No update vertices to build vector index.") context.update( @@ -156,10 +85,4 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: "added_vid_vector_num": len(added_vids), } ) -======= - self.vid_index.save_index_by_name(huge_settings.graph_name, "graph_vids") - else: - log.debug("No update vertices to build vector index.") - context.update({"removed_vid_vector_num": removed_num, "added_vid_vector_num": len(added_vids)}) ->>>>>>> 38dce0b (feat(llm): vector db finished) return context diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index 8d1f3394f..b055cf62e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -59,23 +59,19 @@ def __init__( self.embedding.get_embedding_dim(), "gremlin_examples" ) - def _get_match_result( - self, context: Dict[str, Any], query: str - ) -> List[Dict[str, Any]]: + def _get_match_result(self, context: Dict[str, Any], query: str) -> List[Dict[str, Any]]: if self.num_examples <= 0: return [] query_embedding = context.get("query_embedding") if not isinstance(query_embedding, list): query_embedding = self.embedding.get_texts_embeddings([query])[0] - return self.vector_index.search( - query_embedding, self.num_examples, dis_threshold=1.8 - ) + return self.vector_index.search(query_embedding, self.num_examples, dis_threshold=1.8) def _build_default_example_index(self): - properties = pd.read_csv( - os.path.join(resource_path, "demo", "text2gremlin.csv") - ).to_dict(orient="records") + properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict( + orient="records" + ) from concurrent.futures import ThreadPoolExecutor # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py index c952e5232..49e712d01 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py @@ -19,22 +19,11 @@ import os from typing import Any, Dict, List, Literal, Tuple -<<<<<<< HEAD -<<<<<<< HEAD -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index import VectorIndex -======= -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) -======= from pyhugegraph.client import PyHugeClient from hugegraph_llm.config import huge_settings, resource_path from hugegraph_llm.indices.vector_index.base import VectorStoreBase ->>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import get_filename_prefix, get_index_folder_name from hugegraph_llm.utils.log import log @@ -50,38 +39,10 @@ def __init__( topk_per_keyword: int = huge_settings.topk_per_keyword, vector_dis_threshold: float = huge_settings.vector_dis_threshold, ): -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(embedding, "model_name", None) - ) - self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) -======= - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(embedding, "model_name", None) - ) - self.vector_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) -======= self.index_dir = str(os.path.join(resource_path, huge_settings.graph_name, "graph_vids")) -======= - self.index_dir = str( - os.path.join(resource_path, huge_settings.graph_name, "graph_vids") - ) ->>>>>>> 3aeef7d (fix) self.vector_index = vector_index.from_name( embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids" ) ->>>>>>> 38dce0b (feat(llm): vector db finished) self.embedding = embedding self.by = by self.topk_per_query = topk_per_query @@ -104,9 +65,7 @@ def _exact_match_vids(self, keywords: List[str]) -> Tuple[List[str], List[str]]: possible_vids.update([f"{i + 1}:{keyword}" for keyword in keywords]) vids_str = ",".join([f"'{vid}'" for vid in possible_vids]) - resp = self._client.gremlin().exec( - SemanticIdQuery.ID_QUERY_TEMPL.format(vids_str=vids_str) - ) + resp = self._client.gremlin().exec(SemanticIdQuery.ID_QUERY_TEMPL.format(vids_str=vids_str)) searched_vids = [v["id"] for v in resp["data"]] unsearched_keywords = set(keywords) @@ -120,17 +79,11 @@ def _exact_match_vids(self, keywords: List[str]) -> Tuple[List[str], List[str]]: def _fuzzy_match_vids(self, keywords: List[str]) -> List[str]: fuzzy_match_result = [] for keyword in keywords: -<<<<<<< HEAD - keyword_vector = self.embedding.get_texts_embeddings([keyword])[0] + keyword_vector = self.embedding.get_text_embedding(keyword) results = self.vector_index.search( keyword_vector, top_k=self.topk_per_keyword, dis_threshold=float(self.vector_dis_threshold), -======= - keyword_vector = self.embedding.get_texts_embeddings([keyword])[0] - results = self.vector_index.search( - keyword_vector, top_k=self.topk_per_keyword, dis_threshold=float(self.vector_dis_threshold) ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) ) if results: fuzzy_match_result.extend(results[: self.topk_per_keyword]) @@ -140,7 +93,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: graph_query_list = set() if self.by == "query": query = context["query"] - query_vector = self.embedding.get_texts_embeddings([query])[0] + query_vector = self.embedding.get_text_embedding(query) results = self.vector_index.search(query_vector, top_k=self.topk_per_query) if results: graph_query_list.update(results[: self.topk_per_query]) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py index c83fa7781..0443b15c8 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py @@ -61,9 +61,7 @@ def __init__(self, llm: BaseLLM, embedding: Optional[BaseEmbedding] = None, grap self.graph = graph self.result = None - def import_schema( - self, from_hugegraph=None, from_extraction=None, from_user_defined=None - ): + def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): if from_hugegraph: self.operators.append(SchemaManager(from_hugegraph)) elif from_user_defined: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py index b0864816d..2140abca2 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py @@ -62,19 +62,9 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: -<<<<<<< HEAD -<<<<<<< HEAD - context_str = f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") -======= - context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip("\n") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= - context_str = ( - f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( - "\n" - ) + context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( + "\n" ) ->>>>>>> 3aeef7d (fix) final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -84,7 +74,13 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: graph_result_context, vector_result_context = self.handle_vector_graph(context) context = asyncio.run( - self.async_generate(context, context_head_str, context_tail_str, vector_result_context, graph_result_context) + self.async_generate( + context, + context_head_str, + context_tail_str, + vector_result_context, + graph_result_context, + ) ) return context @@ -94,12 +90,8 @@ def init_llm(self, context): if self._question is None: self._question = context.get("query") or None assert self._question is not None, "No question for synthesizing." - context_head_str = ( - context.get("synthesize_context_head") or self._context_head or "" - ) - context_tail_str = ( - context.get("synthesize_context_tail") or self._context_tail or "" - ) + context_head_str = context.get("synthesize_context_head") or self._context_head or "" + context_tail_str = context.get("synthesize_context_tail") or self._context_tail or "" return context_head_str, context_tail_str def handle_vector_graph(self, context): @@ -123,25 +115,13 @@ def handle_vector_graph(self, context): log.warning(graph_result_context) return graph_result_context, vector_result_context - async def run_streaming( - self, context: Dict[str, Any] - ) -> AsyncGenerator[Dict[str, Any], None]: + async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[str, Any], None]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: -<<<<<<< HEAD -<<<<<<< HEAD - context_str = f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") -======= - context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip("\n") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= - context_str = ( - f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( - "\n" - ) + context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( + "\n" ) ->>>>>>> 3aeef7d (fix) final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -153,7 +133,11 @@ async def run_streaming( graph_result_context, vector_result_context = self.handle_vector_graph(context) async for context in self.async_streaming_generate( - context, context_head_str, context_tail_str, vector_result_context, graph_result_context + context, + context_head_str, + context_tail_str, + vector_result_context, + graph_result_context, ): yield context @@ -169,21 +153,11 @@ async def async_generate( async_tasks = {} if self._raw_answer: final_prompt = self._question - async_tasks["raw_task"] = asyncio.create_task( - self._llm.agenerate(prompt=final_prompt) - ) + async_tasks["raw_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) if self._vector_only_answer: -<<<<<<< HEAD -<<<<<<< HEAD - context_str = f"{context_head_str}\n" f"{vector_result_context}\n" f"{context_tail_str}".strip("\n") -======= - context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip("\n") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip( "\n" ) ->>>>>>> 3aeef7d (fix) final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -192,19 +166,9 @@ async def async_generate( self._llm.agenerate(prompt=final_prompt) ) if self._graph_only_answer: -<<<<<<< HEAD -<<<<<<< HEAD - context_str = f"{context_head_str}\n" f"{graph_result_context}\n" f"{context_tail_str}".strip("\n") -======= - context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip("\n") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= - context_str = ( - f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( - "\n" - ) + context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( + "\n" ) ->>>>>>> 3aeef7d (fix) final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -216,19 +180,7 @@ async def async_generate( context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" -<<<<<<< HEAD -<<<<<<< HEAD - context_str = f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") -======= context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= - context_str = ( - f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip( - "\n" - ) - ) ->>>>>>> 3aeef7d (fix) final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -250,7 +202,14 @@ async def async_generate( context[context_key] = response log.debug("Query Answer: %s", response) - ops = sum([self._raw_answer, self._vector_only_answer, self._graph_only_answer, self._graph_vector_answer]) + ops = sum( + [ + self._raw_answer, + self._vector_only_answer, + self._graph_only_answer, + self._graph_vector_answer, + ] + ) context["call_count"] = context.get("call_count", 0) + ops return context @@ -274,41 +233,25 @@ async def async_streaming_generate( ) auto_id += 1 if self._vector_only_answer: -<<<<<<< HEAD -<<<<<<< HEAD - context_str = f"{context_head_str}\n" f"{vector_result_context}\n" f"{context_tail_str}".strip("\n") -======= - context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip("\n") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip( "\n" ) ->>>>>>> 3aeef7d (fix) final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question ) async_generators.append( self.__llm_generate_with_meta_info( - task_id=auto_id, target_key="vector_only_answer", prompt=final_prompt + task_id=auto_id, + target_key="vector_only_answer", + prompt=final_prompt, ) ) auto_id += 1 if self._graph_only_answer: -<<<<<<< HEAD -<<<<<<< HEAD - context_str = f"{context_head_str}\n" f"{graph_result_context}\n" f"{context_tail_str}".strip("\n") -======= - context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip("\n") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= - context_str = ( - f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( - "\n" - ) + context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( + "\n" ) ->>>>>>> 3aeef7d (fix) final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -323,38 +266,33 @@ async def async_streaming_generate( context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" -<<<<<<< HEAD -<<<<<<< HEAD - context_str = f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") -======= context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= - context_str = ( - f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip( - "\n" - ) - ) ->>>>>>> 3aeef7d (fix) final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question ) async_generators.append( self.__llm_generate_with_meta_info( - task_id=auto_id, target_key="graph_vector_answer", prompt=final_prompt + task_id=auto_id, + target_key="graph_vector_answer", + prompt=final_prompt, ) ) auto_id += 1 - ops = sum([self._raw_answer, self._vector_only_answer, self._graph_only_answer, self._graph_vector_answer]) + ops = sum( + [ + self._raw_answer, + self._vector_only_answer, + self._graph_only_answer, + self._graph_vector_answer, + ] + ) context["call_count"] = context.get("call_count", 0) + ops async_tasks = [asyncio.create_task(anext(gen)) for gen in async_generators] while True: - done, _ = await asyncio.wait( - async_tasks, return_when=asyncio.FIRST_COMPLETED - ) + done, _ = await asyncio.wait(async_tasks, return_when=asyncio.FIRST_COMPLETED) stop_task_num = 0 for task in done: try: @@ -368,9 +306,7 @@ async def async_streaming_generate( break yield context - async def __llm_generate_with_meta_info( - self, task_id: int, target_key: str, prompt: str - ): + async def __llm_generate_with_meta_info(self, task_id: int, target_key: str, prompt: str): # FIXME: Expected type 'AsyncIterable', got 'Coroutine[Any, Any, AsyncGenerator[str, None]]' instead async for token in self._llm.agenerate_streaming(prompt=prompt): yield task_id, target_key, token diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index 2c0244d57..fd4583263 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -48,9 +48,7 @@ def _extract_response(self, response: str, label: str = "gremlin") -> str: return match.group(1).strip() return response.strip() - def _format_examples( - self, examples: Optional[List[Dict[str, str]]] - ) -> Optional[str]: + def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional[str]: if not examples: return None example_strings = [] @@ -88,9 +86,7 @@ def _format_properties(self, properties: Optional[List[tuple]]) -> Optional[str] async def async_generate(self, context: Dict[str, Any]): async_tasks = {} query = context.get("query") - raw_example = [ - {"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"} - ] + raw_example = [{"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"}] raw_prompt = self.gremlin_prompt.format( query=query, schema=self.schema, @@ -98,9 +94,7 @@ async def async_generate(self, context: Dict[str, Any]): vertices=self._format_vertices(vertices=self.vertices), properties=self._format_properties(properties=None), ) - async_tasks["raw_answer"] = asyncio.create_task( - self.llm.agenerate(prompt=raw_prompt) - ) + async_tasks["raw_answer"] = asyncio.create_task(self.llm.agenerate(prompt=raw_prompt)) examples = context.get("match_result") init_prompt = self.gremlin_prompt.format( @@ -130,9 +124,7 @@ async def async_generate(self, context: Dict[str, Any]): def sync_generate(self, context: Dict[str, Any]): query = context.get("query") - raw_example = [ - {"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"} - ] + raw_example = [{"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"}] raw_prompt = self.gremlin_prompt.format( query=query, schema=self.schema, 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 00a237077..a9ac4c050 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 @@ -16,7 +16,7 @@ # under the License. import re -from typing import List, Any, Dict, Optional +from typing import Any, Dict, List, Optional from hugegraph_llm.document.chunk_split import ChunkSplitter from hugegraph_llm.models.llms.base import BaseLLM @@ -27,11 +27,11 @@ ## Basic Rules 1. The output format must be: (X,Y,Z) - LABEL -In this format, Y must be a value from "properties" or "edge_label", +In this format, Y must be a value from "properties" or "edge_label", and LABEL must be X's vertex_label or Y's edge_label. 2. Don't extract attribute/property fields that do not exist in the given schema 3. Ensure the extract property is in the same type as the schema (like 'age' should be a number) -4. Translate the given schema filed into Chinese if the given text is Chinese but the schema is in English (Optional) +4. Translate the given schema filed into Chinese if the given text is Chinese but the schema is in English (Optional) ## Example (Note: Update the example to correspond to the given text and schema) ### Input example: @@ -75,24 +75,9 @@ def generate_extract_triple_prompt(text, schema=None) -> str: if schema: return schema_real_prompt -<<<<<<< HEAD -<<<<<<< HEAD - log.warning( -<<<<<<< HEAD - "Recommend to provide a graph schema to improve the extraction accuracy. " - "Now using the default schema." -======= - "Recommend to provide a graph schema to improve the extraction accuracy. " "Now using the default schema." ->>>>>>> 87ee5d3 (style: format code with black line-length 120) - ) -======= - log.warning("Recommend to provide a graph schema to improve the extraction accuracy. Now using the default schema.") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= log.warning( "Recommend to provide a graph schema to improve the extraction accuracy. Now using the default schema." ) ->>>>>>> 3aeef7d (fix) return text_based_prompt @@ -121,23 +106,9 @@ def extract_triples_by_regex_with_schema(schema, text, graph): # TODO: use a more efficient way to compare the extract & input property p_lower = p.lower() for vertex in schema["vertices"]: -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - if vertex["vertex_label"] == label and any( - pp.lower() == p_lower for pp in vertex["properties"] - ): -======= - if vertex["vertex_label"] == label and any(pp.lower() == p_lower for pp in vertex["properties"]): ->>>>>>> 87ee5d3 (style: format code with black line-length 120) -======= - if vertex["vertex_label"] == label and any(pp.lower() == p_lower for pp in vertex["properties"]): ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= if vertex["vertex_label"] == label and any( pp.lower() == p_lower for pp in vertex["properties"] ): ->>>>>>> 3aeef7d (fix) id = f"{label}-{s}" if id not in vertices_dict: vertices_dict[id] = { @@ -154,7 +125,6 @@ def extract_triples_by_regex_with_schema(schema, text, graph): source_label = edge["source_vertex_label"] source_id = f"{source_label}-{s}" if source_id not in vertices_dict: -<<<<<<< HEAD vertices_dict[source_id] = { "id": source_id, "name": s, @@ -178,16 +148,8 @@ def extract_triples_by_regex_with_schema(schema, text, graph): "properties": {}, } ) -======= - vertices_dict[source_id] = {"id": source_id, "name": s, "label": source_label, "properties": {}} - target_label = edge["target_vertex_label"] - target_id = f"{target_label}-{o}" - if target_id not in vertices_dict: - vertices_dict[target_id] = {"id": target_id, "name": o, "label": target_label, "properties": {}} - graph["edges"].append({"start": source_id, "end": target_id, "type": label, "properties": {}}) ->>>>>>> 87ee5d3 (style: format code with black line-length 120) break - graph["vertices"] = list(vertices_dict.values()) + graph["vertices"] = vertices_dict.values() class InfoExtract: @@ -235,27 +197,8 @@ def valid(self, element_id: str, max_length: int = 256) -> bool: return True def _filter_long_id(self, graph) -> Dict[str, List[Any]]: -<<<<<<< HEAD graph["vertices"] = [vertex for vertex in graph["vertices"] if self.valid(vertex["id"])] -<<<<<<< HEAD -<<<<<<< HEAD graph["edges"] = [ edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"]) ] -======= - graph["edges"] = [edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"])] ->>>>>>> 87ee5d3 (style: format code with black line-length 120) -======= - graph["edges"] = [edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"])] ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= - graph["vertices"] = [ - vertex for vertex in graph["vertices"] if self.valid(vertex["id"]) - ] - graph["edges"] = [ - edge - for edge in graph["edges"] - if self.valid(edge["start"]) and self.valid(edge["end"]) - ] ->>>>>>> 3aeef7d (fix) return graph diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index 420fc9776..0fa1c0f04 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -113,8 +113,6 @@ def _extract_keywords_from_response( sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: results.update( - w - for w in sub_tokens - if w not in NLTKHelper().stopwords(lang=self._language) + w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language) ) return results diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py index a45812393..058d1bce9 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py @@ -31,9 +31,7 @@ def __init__(self, llm: BaseLLM): def _load_few_shot_example(self, example_name: str) -> Dict[str, Any]: """Loads and finds the specified few-shot example from the unified JSON file.""" - examples_path = os.path.join( - resource_path, "prompt_examples", "prompt_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") if not os.path.exists(examples_path): raise FileNotFoundError(f"Examples file not found: {examples_path}") with open(examples_path, "r", encoding="utf-8") as f: @@ -41,9 +39,7 @@ def _load_few_shot_example(self, example_name: str) -> Dict[str, Any]: for example in all_examples: if example.get("name") == example_name: return example - raise ValueError( - f"Example with name '{example_name}' not found in prompt_examples.json" - ) + raise ValueError(f"Example with name '{example_name}' not found in prompt_examples.json") def run(self, context: Dict[str, Any]) -> Dict[str, Any]: """Executes the core logic of prompt generation.""" @@ -52,9 +48,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: example_name = context.get("example_name") if not all([source_text, scenario, example_name]): - raise ValueError( - "Missing required context: source_text, scenario, or example_name." - ) + raise ValueError("Missing required context: source_text, scenario, or example_name.") few_shot_example = self._load_few_shot_example(example_name) meta_prompt = prompt_tpl.generate_extract_prompt_template.format( 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 3ba178c88..d517db8b9 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 @@ -67,9 +67,9 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: item_type = item["type"] if item_type == "vertex": label = item["label"] - non_nullable_keys = set( - properties_map[item_type][label]["properties"] - ).difference(set(properties_map[item_type][label]["nullable_keys"])) + non_nullable_keys = set(properties_map[item_type][label]["properties"]).difference( + set(properties_map[item_type][label]["nullable_keys"]) + ) for key in non_nullable_keys: if key not in item["properties"]: item["properties"][key] = "NULL" @@ -82,9 +82,7 @@ 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) -> None: self.llm = llm self.example_prompt = example_prompt self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name @@ -150,9 +148,7 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: and "vertices" in property_graph and "edges" in property_graph ): - log.critical( - "Invalid property graph format; expecting 'vertices' and 'edges'." - ) + log.critical("Invalid property graph format; expecting 'vertices' and 'edges'.") return items # Create sets for valid vertex and edge labels based on the schema @@ -162,9 +158,7 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: def process_items(item_list, valid_labels, item_type): for item in item_list: if not isinstance(item, dict): - log.warning( - "Invalid property graph item type '%s'.", type(item) - ) + log.warning("Invalid property graph item type '%s'.", type(item)) continue if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): log.warning("Invalid item keys '%s'.", item.keys()) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py index ae445c206..1e33514ca 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py @@ -117,13 +117,9 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: raise ValueError("Context must be a dictionary") if "raw_texts" not in context or not isinstance(context["raw_texts"], list): raise ValueError("'raw_texts' must be a list[str]") - if "query_examples" not in context or not isinstance( - context["query_examples"], list - ): + if "query_examples" not in context or not isinstance(context["query_examples"], list): raise ValueError("'query_examples' must be a list[str]") - if "few_shot_schema" not in context or not isinstance( - context["few_shot_schema"], dict - ): + if "few_shot_schema" not in context or not isinstance(context["few_shot_schema"], dict): raise ValueError("'few_shot_schema' must be a dict") raw_texts = context["raw_texts"] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py index 98cc97ccf..6beeb0291 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py @@ -104,9 +104,7 @@ def nodes_schemas_text_to_list_of_dict(nodes_schemas): properties = json.loads(properties) except json.decoder.JSONDecodeError: properties = {} - result.append( - {"label": label, "primary_key": primary_key, "properties": properties} - ) + result.append({"label": label, "primary_key": primary_key, "properties": properties}) return result @@ -118,9 +116,7 @@ def relationships_schemas_text_to_list_of_dict(relationships_schemas): continue start = relationships_schema_list[0].strip().replace('"', "") end = relationships_schema_list[2].strip().replace('"', "") - relationships_schema_type = ( - relationships_schema_list[1].strip().replace('"', "") - ) + relationships_schema_type = relationships_schema_list[1].strip().replace('"', "") properties = re.search(JSON_REGEX, relationships_schema) if properties is None: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py index b5232d268..2914c4b28 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py @@ -23,9 +23,7 @@ from hugegraph_llm.utils.log import log -def log_elapsed_time( - start_time: float, func: Callable, args: tuple, msg: Optional[str] -): +def log_elapsed_time(start_time: float, func: Callable, args: tuple, msg: Optional[str]): elapse_time = time.perf_counter() - start_time unit = "s" if elapse_time < 1: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py index ace3d4b6a..b2f485cea 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py @@ -32,9 +32,7 @@ async def _get_batch_with_progress( return result -async def get_embeddings_parallel( - embedding: BaseEmbedding, vids: list[str] -) -> list[Any]: +async def get_embeddings_parallel(embedding: BaseEmbedding, vids: list[str]) -> list[Any]: """Get embeddings for texts in parallel. This function processes text embeddings asynchronously in parallel, using batching and semaphore @@ -62,9 +60,7 @@ async def get_embeddings_parallel( embeddings = [] with tqdm(total=len(vid_batches)) as pbar: # Create tasks for each batch with progress bar updates - tasks = [ - _get_batch_with_progress(embedding, batch, pbar) for batch in vid_batches - ] + tasks = [_get_batch_with_progress(embedding, batch, pbar) for batch in vid_batches] # Use asyncio.gather() to preserve order batch_results = await asyncio.gather(*tasks) @@ -78,9 +74,7 @@ async def get_embeddings_parallel( def get_filename_prefix(embedding_type: str = None, model_name: str = None) -> str: """Generate filename based on model name.""" - if not ( - model_name and model_name.strip() and embedding_type and embedding_type.strip() - ): + if not (model_name and model_name.strip() and embedding_type and embedding_type.strip()): return "" # Sanitize model_name to prevent path traversal or invalid filename chars safe_embedding_type = embedding_type.replace("/", "_").replace("\\", "_").strip() diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 6fab74bf9..2f48af332 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -22,14 +22,8 @@ from typing import Any, Dict, Optional, Union import gradio as gr -from hugegraph_llm.flows.scheduler import SchedulerSingleton -from .embedding_utils import get_filename_prefix, get_index_folder_name -from .hugegraph_utils import get_hg_client, clean_hg_data -from .log import log -from .vector_index_utils import read_documents -from ..config import resource_path, huge_settings, llm_settings -from ..indices.vector_index.faiss_vector_store import FaissVectorIndex +from ..config import huge_settings, index_settings, resource_path from ..models.embeddings.init_embedding import Embeddings from ..models.llms.init_llm import LLMs from ..operators.kg_construction_task import KgBuilder @@ -39,19 +33,13 @@ def get_graph_index_info(): - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) graph_summary_info = builder.fetch_graph_data().run() - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None) - ) - vector_index = FaissVectorIndex.from_index_file( - str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix, record_miss=False - ) + vector_index = get_vector_index_class(index_settings.cur_vector_index) vector_index_entity = vector_index.from_name( - Embeddings().get_embedding().get_embedding_dim(), huge_settings.graph_name, "chunks" + Embeddings().get_embedding().get_embedding_dim(), + huge_settings.graph_name, + "chunks", ) vector_index_info = vector_index_entity.get_vector_index_info() graph_summary_info["vid_index"] = { @@ -98,20 +86,16 @@ def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: return None -def extract_graph_origin(input_file, input_text, schema, example_prompt) -> str: +def extract_graph(input_file, input_text, schema, example_prompt) -> str: texts = read_documents(input_file, input_text) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) if not schema: return "ERROR: please input with correct schema/format." error_message = parse_schema(schema, builder) if error_message: return error_message - builder.chunk_split(texts, "document", "zh").extract_info( - example_prompt, "property_graph" - ) + builder.chunk_split(texts, "document", "zh").extract_info(example_prompt, "property_graph") try: context = builder.run() @@ -136,38 +120,16 @@ def extract_graph_origin(input_file, input_text, schema, example_prompt) -> str: raise gr.Error(str(e)) -def extract_graph(input_file, input_text, schema, example_prompt) -> str: - texts = read_documents(input_file, input_text) - scheduler = SchedulerSingleton.get_instance() - if not schema: - return "ERROR: please input with correct schema/format." - - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) - if not schema: - return "ERROR: please input with correct schema/format." - - error_message = parse_schema(schema, builder) - if error_message: - return error_message - builder.chunk_split(texts, "document", "zh").extract_info( - example_prompt, "property_graph" - ) - - try: - return scheduler.schedule_flow( - "graph_extract", schema, texts, example_prompt, "property_graph" - ) - except Exception as e: # pylint: disable=broad-exception-caught - log.error(e) - raise gr.Error(str(e)) - - def update_vid_embedding(): - scheduler = SchedulerSingleton.get_instance() + vector_index = get_vector_index_class(index_settings.cur_vector_index) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder.fetch_graph_data().build_vertex_id_semantic_index(vector_index) + log.debug("Operators: %s", builder.operators) try: - return scheduler.schedule_flow("update_vid_embeddings") + context = builder.run() + removed_num = context["removed_vid_vector_num"] + added_num = context["added_vid_vector_num"] + return f"Removed {removed_num} vectors, added {added_num} vectors." except Exception as e: # pylint: disable=broad-exception-caught log.error(e) raise gr.Error(str(e)) @@ -177,9 +139,7 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: data_json = json.loads(data.strip()) log.debug("Import graph data: %s", data) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) if schema: error_message = parse_schema(schema, builder) if error_message: @@ -198,10 +158,42 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: def build_schema(input_text, query_example, few_shot): - scheduler = SchedulerSingleton.get_instance() + context = { + "raw_texts": [input_text] if input_text else [], + "query_examples": [], + "few_shot_schema": {}, + } + + if few_shot: + try: + context["few_shot_schema"] = json.loads(few_shot) + except json.JSONDecodeError as e: + raise gr.Error(f"Few Shot Schema is not in a valid JSON format: {e}") from e + + if query_example: + try: + parsed_examples = json.loads(query_example) + # Validate and retain the description and gremlin fields + context["query_examples"] = [ + { + "description": ex.get("description", ""), + "gremlin": ex.get("gremlin", ""), + } + for ex in parsed_examples + if isinstance(ex, dict) and "description" in ex and "gremlin" in ex + ] + except json.JSONDecodeError as e: + raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e + + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) try: - return scheduler.schedule_flow( - "build_schema", input_text, query_example, few_shot - ) + schema = builder.build_schema().run(context) + except Exception as e: + log.error("Failed to generate schema: %s", e) + raise gr.Error(f"Schema generation failed: {e}") from e + try: + formatted_schema = json.dumps(schema, ensure_ascii=False, indent=2) + return formatted_schema except (TypeError, ValueError) as e: - raise gr.Error(f"Schema generation failed: {e}") + log.error("Failed to format schema: %s", e) + return str(schema) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index 082d6bb8c..6da7a6567 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -29,9 +29,7 @@ MAX_BACKUP_DIRS = 7 MAX_VERTICES = 100000 MAX_EDGES = 200000 -BACKUP_DIR = str( - os.path.join(resource_path, "backup-graph-data-4020", huge_settings.graph_name) -) +BACKUP_DIR = str(os.path.join(resource_path, "backup-graph-data-4020", huge_settings.graph_name)) def run_gremlin_query(query, fmt=True): @@ -58,33 +56,21 @@ def init_hg_test_data(): schema.vertexLabel("Person").properties( "name", "birthDate" ).useCustomizeStringId().ifNotExist().create() - schema.vertexLabel("Movie").properties( - "name" - ).useCustomizeStringId().ifNotExist().create() - schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel( - "Movie" - ).ifNotExist().create() + schema.vertexLabel("Movie").properties("name").useCustomizeStringId().ifNotExist().create() + schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() - schema.indexLabel("PersonByName").onV("Person").by( - "name" - ).secondary().ifNotExist().create() - schema.indexLabel("MovieByName").onV("Movie").by( - "name" - ).secondary().ifNotExist().create() + schema.indexLabel("PersonByName").onV("Person").by("name").secondary().ifNotExist().create() + schema.indexLabel("MovieByName").onV("Movie").by("name").secondary().ifNotExist().create() graph = client.graph() - graph.addVertex( - "Person", {"name": "Al Pacino", "birthDate": "1940-04-25"}, id="Al Pacino" - ) + graph.addVertex("Person", {"name": "Al Pacino", "birthDate": "1940-04-25"}, id="Al Pacino") graph.addVertex( "Person", {"name": "Robert De Niro", "birthDate": "1943-08-17"}, id="Robert De Niro", ) graph.addVertex("Movie", {"name": "The Godfather"}, id="The Godfather") - graph.addVertex( - "Movie", {"name": "The Godfather Part II"}, id="The Godfather Part II" - ) + graph.addVertex("Movie", {"name": "The Godfather Part II"}, id="The Godfather Part II") graph.addVertex( "Movie", {"name": "The Godfather Coda The Death of Michael Corleone"}, @@ -93,9 +79,7 @@ def init_hg_test_data(): graph.addEdge("ActedIn", "Al Pacino", "The Godfather", {}) graph.addEdge("ActedIn", "Al Pacino", "The Godfather Part II", {}) - graph.addEdge( - "ActedIn", "Al Pacino", "The Godfather Coda The Death of Michael Corleone", {} - ) + graph.addEdge("ActedIn", "Al Pacino", "The Godfather Coda The Death of Michael Corleone", {}) graph.addEdge("ActedIn", "Robert De Niro", "The Godfather Part II", {}) schema.getSchema() graph.close() @@ -134,9 +118,7 @@ def backup_data(): } vertexlabels = client.schema().getSchema()["vertexlabels"] - all_pk_flag = all( - data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels - ) + all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) for filename, query in files.items(): write_backup_file(client, backup_subdir, filename, query, all_pk_flag) @@ -216,9 +198,7 @@ def manage_backup_retention(): # TODO: In the path demo/rag_demo/configs_block.py, # there is a function test_api_connection that is similar to this function, # but it is not straightforward to reuse -def check_graph_db_connection( - url: str, name: str, user: str, pwd: str, graph_space: str -) -> bool: +def check_graph_db_connection(url: str, name: str, user: str, pwd: str, graph_space: str) -> bool: try: if graph_space and graph_space.strip(): test_url = f"{url}/graphspaces/{graph_space}/graphs/{name}/schema" diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index df48ea1e9..4f56db636 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -51,9 +51,7 @@ def read_documents(input_file, input_text): texts.append(text) elif full_path.endswith(".pdf"): # TODO: support PDF file - raise gr.Error( - "PDF will be supported later! Try to upload text/docx now" - ) + raise gr.Error("PDF will be supported later! Try to upload text/docx now") else: raise gr.Error("Please input txt or docx file.") else: @@ -89,14 +87,8 @@ def build_vector_index(input_file, input_text): if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) - context = ( - builder.chunk_split(texts, "paragraph", "zh") - .build_vector_index(vector_index) - .run() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + context = builder.chunk_split(texts, "paragraph", "zh").build_vector_index(vector_index).run() return json.dumps(context, ensure_ascii=False, indent=2) From 6daf82cae37105924441bba9753ac9b4448bccb5 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Wed, 3 Sep 2025 15:54:32 +0800 Subject: [PATCH 27/71] fix schema g & prompt g --- .../demo/rag_demo/configs_block.py | 57 +-- .../demo/rag_demo/vector_graph_block.py | 365 +++--------------- .../operators/index_op/build_vector_index.py | 49 +-- .../operators/kg_construction_task.py | 14 +- 4 files changed, 87 insertions(+), 398 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index d4415700b..112892e7c 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -33,7 +33,7 @@ current_llm = "chat" -def test_litellm_embedding(api_key, api_base, model_name, model_dim) -> int: +def test_litellm_embedding(api_key, api_base, model_name) -> int: llm_client = LiteLLMEmbedding( api_key=api_key, api_base=api_base, @@ -125,15 +125,13 @@ def config_qianfan_model(arg1, arg2, arg3=None, settings_prefix=None, origin_cal return status_code -def apply_embedding_config(arg1, arg2, arg3, arg4, origin_call=None) -> int: +def apply_embedding_config(arg1, arg2, arg3, origin_call=None) -> int: status_code = -1 embedding_option = llm_settings.embedding_type - arg4 = int(arg4) if embedding_option == "openai": llm_settings.openai_embedding_api_key = arg1 llm_settings.openai_embedding_api_base = arg2 llm_settings.openai_embedding_model = arg3 - llm_settings.openai_embedding_model_dim = arg4 test_url = llm_settings.openai_embedding_api_base + "/embeddings" headers = {"Authorization": f"Bearer {arg1}"} data = {"model": arg3, "input": "test"} @@ -144,14 +142,12 @@ def apply_embedding_config(arg1, arg2, arg3, arg4, origin_call=None) -> int: llm_settings.ollama_embedding_host = arg1 llm_settings.ollama_embedding_port = int(arg2) llm_settings.ollama_embedding_model = arg3 - llm_settings.ollama_embedding_model_dim = arg4 status_code = test_api_connection(f"http://{arg1}:{arg2}", origin_call=origin_call) elif embedding_option == "litellm": llm_settings.litellm_embedding_api_key = arg1 llm_settings.litellm_embedding_api_base = arg2 llm_settings.litellm_embedding_model = arg3 - llm_settings.litellm_embedding_model_dim = arg4 - status_code = test_litellm_embedding(arg1, arg2, arg3, arg4) + status_code = test_litellm_embedding(arg1, arg2, arg3) llm_settings.update_env() gr.Info("Configured!") return status_code @@ -168,10 +164,10 @@ def apply_reranker_config( if reranker_option == "cohere": llm_settings.reranker_api_key = reranker_api_key llm_settings.reranker_model = reranker_model - llm_settings.cohere_base_url = cohere_base_url # type:ignore + llm_settings.cohere_base_url = cohere_base_url headers = {"Authorization": f"Bearer {reranker_api_key}"} status_code = test_api_connection( - cohere_base_url.rsplit("/", 1)[0] + "/check-api-key", # type:ignore + cohere_base_url.rsplit("/", 1)[0] + "/check-api-key", method="POST", headers=headers, origin_call=origin_call, @@ -530,22 +526,22 @@ def text2gql_llm_settings(llm_type): elif llm_type == "litellm": llm_config_input = [ gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_text2gql_api_key"), + value=getattr(llm_settings, "litellm_text2gql_api_key"), label="api_key", type="password", ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_text2gql_api_base"), + value=getattr(llm_settings, "litellm_text2gql_api_base"), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_text2gql_language_model"), + value=getattr(llm_settings, "litellm_text2gql_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_text2gql_tokens"), + value=getattr(llm_settings, "litellm_text2gql_tokens"), label="max_token", ), ] @@ -568,78 +564,63 @@ def embedding_settings(embedding_type): with gr.Row(): embedding_config_input = [ gr.Textbox( - value=lambda: llm_settings.openai_embedding_api_key, + value=llm_settings.openai_embedding_api_key, label="api_key", type="password", ), gr.Textbox( - value=lambda: llm_settings.openai_embedding_api_base, + value=llm_settings.openai_embedding_api_base, label="api_base", ), gr.Textbox( - value=lambda: llm_settings.openai_embedding_model, + value=llm_settings.openai_embedding_model, label="model_name", ), - gr.Textbox( - value=lambda: str(llm_settings.openai_embedding_model_dim), - label="model_dim", - ), ] elif embedding_type == "ollama/local": with gr.Row(): embedding_config_input = [ gr.Textbox( - value=lambda: llm_settings.ollama_embedding_host, + value=llm_settings.ollama_embedding_host, label="host", ), gr.Textbox( - value=lambda: str(llm_settings.ollama_embedding_port), + value=str(llm_settings.ollama_embedding_port), label="port", ), gr.Textbox( - value=lambda: llm_settings.ollama_embedding_model, + value=llm_settings.ollama_embedding_model, label="model_name", ), - gr.Textbox( - value=lambda: str(llm_settings.ollama_embedding_model_dim), - label="model_dim", - ), ] elif embedding_type == "litellm": with gr.Row(): embedding_config_input = [ gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_embedding_api_key"), + value=getattr(llm_settings, "litellm_embedding_api_key"), label="api_key", type="password", ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_embedding_api_base"), + value=getattr(llm_settings, "litellm_embedding_api_base"), label="api_base", info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_embedding_model"), + value=getattr(llm_settings, "litellm_embedding_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/embedding/supported_embedding", ), - gr.Textbox( - value=lambda: getattr(llm_settings, "litellm_embedding_model_dim"), - label="model_dim", - type="text", - ), ] else: embedding_config_input = [ gr.Textbox(value="", visible=False), gr.Textbox(value="", visible=False), gr.Textbox(value="", visible=False), - gr.Textbox(value="", visible=False), ] embedding_config_button = gr.Button("Apply Configuration") - # Call the separate apply_embedding_configuration function here embedding_config_button.click( # pylint: disable=no-member fn=apply_embedding_config, inputs=embedding_config_input, # pylint: disable=no-member @@ -715,7 +696,7 @@ def get_header_with_language_indicator(language: str) -> str: language_class = language.lower() if language == "CN": - title_text = "当前 prompt 语言：中文 (CN)" + title_text = "当前prompt语言: 中文 (CN)" else: title_text = "Current prompt Language: English (EN)" html_content = f""" diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index f011e2082..1e8124b81 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -23,37 +23,24 @@ import gradio as gr -<<<<<<< HEAD -from hugegraph_llm.config import huge_settings -from hugegraph_llm.config import prompt -from hugegraph_llm.config import resource_path -from hugegraph_llm.flows.scheduler import SchedulerSingleton -======= -from hugegraph_llm.config import huge_settings, prompt ->>>>>>> 38dce0b (feat(llm): vector db finished) +from hugegraph_llm.config import huge_settings, prompt, resource_path +from hugegraph_llm.models.llms.init_llm import LLMs +from hugegraph_llm.operators.llm_op.prompt_generate import PromptGenerate from hugegraph_llm.utils.graph_index_utils import ( + build_schema, clean_all_graph_data, clean_all_graph_index, extract_graph, get_graph_index_info, import_graph_data, -<<<<<<< HEAD - build_schema, + update_vid_embedding, ) from hugegraph_llm.utils.hugegraph_utils import check_graph_db_connection from hugegraph_llm.utils.log import log from hugegraph_llm.utils.vector_index_utils import ( - clean_vector_index, build_vector_index, - get_vector_index_info, -) -======= - update_vid_embedding, + clean_vector_index, ) -from hugegraph_llm.utils.hugegraph_utils import check_graph_db_connection -from hugegraph_llm.utils.log import log -from hugegraph_llm.utils.vector_index_utils import build_vector_index, clean_vector_index, get_vector_index_info ->>>>>>> 38dce0b (feat(llm): vector db finished) def store_prompt(doc, schema, example_prompt): @@ -71,17 +58,27 @@ def store_prompt(doc, schema, example_prompt): def generate_prompt_for_ui(source_text, scenario, example_name): """ - Handles the UI logic for generating a new prompt using the new workflow architecture. + Handles the UI logic for generating a new prompt. It calls the PromptGenerate operator. """ if not all([source_text, scenario, example_name]): - gr.Warning("Please provide original text, expected scenario, and select an example!") + gr.Warning( + "Please provide original text, expected scenario, and select an example!" + ) return gr.update() try: - # using new architecture - scheduler = SchedulerSingleton.get_instance() - result = scheduler.schedule_flow("prompt_generate", source_text, scenario, example_name) + prompt_generator = PromptGenerate(llm=LLMs().get_chat_llm()) + context = { + "source_text": source_text, + "scenario": scenario, + "example_name": example_name, + } + result_context = prompt_generator.run(context) + # Presents the result of generating prompt + generated_prompt = result_context.get( + "generated_extract_prompt", "Generation failed. Please check the logs." + ) gr.Info("Prompt generated successfully!") - return result + return generated_prompt except Exception as e: log.error("Error generating Prompt: %s", e, exc_info=True) raise gr.Error(f"Error generating Prompt: {e}") from e @@ -90,7 +87,9 @@ def generate_prompt_for_ui(source_text, scenario, example_name): def load_example_names(): """Load all candidate examples""" try: - examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "prompt_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return [example.get("name", "Unnamed example") for example in examples] @@ -104,37 +103,27 @@ def load_query_examples(): language = getattr( prompt, "language", -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -======= ->>>>>>> 3aeef7d (fix) - ( getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") - else "EN" - ), -<<<<<<< HEAD -======= - getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") else "EN", ->>>>>>> 87ee5d3 (style: format code with black line-length 120) -======= - (getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") else "EN"), ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= ->>>>>>> 3aeef7d (fix) + else "EN", ) if language.upper() == "CN": - examples_path = os.path.join(resource_path, "prompt_examples", "query_examples_CN.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "query_examples_CN.json" + ) else: - examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "query_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) except (FileNotFoundError, json.JSONDecodeError): try: - examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "query_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -145,7 +134,9 @@ def load_query_examples(): def load_schema_fewshot_examples(): """Load few-shot examples from a JSON file""" try: - examples_path = os.path.join(resource_path, "prompt_examples", "schema_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "schema_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -156,10 +147,14 @@ def load_schema_fewshot_examples(): def update_example_preview(example_name): """Update the display content based on the selected example name.""" try: - examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "prompt_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: all_examples = json.load(f) - selected_example = next((ex for ex in all_examples if ex.get("name") == example_name), None) + selected_example = next( + (ex for ex in all_examples if ex.get("name") == example_name), None + ) if selected_example: return ( @@ -187,11 +182,9 @@ def _create_prompt_helper_block(demo, input_text, info_extract_template): few_shot_dropdown = gr.Dropdown( choices=example_names, label="Select a Few-shot example as a reference", - value=( - example_names[0] - if example_names and example_names[0] != "No available examples" - else None - ), + value=example_names[0] + if example_names and example_names[0] != "No available examples" + else None, ) with gr.Accordion("View example details", open=False): example_desc_preview = gr.Markdown(label="Example description") @@ -204,7 +197,9 @@ def _create_prompt_helper_block(demo, input_text, info_extract_template): interactive=False, ) - generate_prompt_btn = gr.Button("🚀 Auto-generate Graph Extract Prompt", variant="primary") + generate_prompt_btn = gr.Button( + "🚀 Auto-generate Graph Extract Prompt", variant="primary" + ) # Bind the change event of the dropdown menu few_shot_dropdown.change( fn=update_example_preview, @@ -251,7 +246,6 @@ def create_vector_graph_block(): # pylint: disable=no-member # pylint: disable=C0301 # pylint: disable=unexpected-keyword-arg -<<<<<<< HEAD with gr.Blocks() as demo: gr.Markdown( """## Build Vector/Graph Index & Extract Knowledge Graph @@ -266,47 +260,10 @@ def create_vector_graph_block(): - Graph Extract Prompt Header: The user-defined prompt of graph extracting - If already exist the graph data, you should click "**Rebuild vid Index**" to update the index """ -======= - gr.Markdown( - """## Build Vector/Graph Index & Extract Knowledge Graph -- Docs: - - text: Build rag index from plain text - - file: Upload file(s) which should be TXT or .docx (Multiple files can be selected together) -- [Schema](https://hugegraph.apache.org/docs/clients/restful-api/schema/): (Accept **2 types**) - - User-defined Schema (JSON format, follow the [template](https://github.com/apache/incubator-hugegraph-ai/blob/aff3bbe25fa91c3414947a196131be812c20ef11/hugegraph-llm/src/hugegraph_llm/config/config_data.py#L125) - to modify it) - - Specify the name of the HugeGraph graph instance, it will automatically get the schema from it (like - **"hugegraph"**) -- Graph Extract Prompt Header: The user-defined prompt of graph extracting -- If already exist the graph data, you should click "**Rebuild vid Index**" to update the index -""" - ) - - with gr.Row(): - with gr.Column(): - with gr.Tab("text") as tab_upload_text: - input_text = gr.Textbox( - value=prompt.doc_input_text, label="Input Doc(s)", lines=20, show_copy_button=True - ) - with gr.Tab("file") as tab_upload_file: - input_file = gr.File( - value=None, - label="Docs (multi-files can be selected together)", - file_count="multiple", - ) - input_schema = gr.Code(value=prompt.graph_schema, label="Graph Schema", language="json", lines=15, max_lines=29) - info_extract_template = gr.Code( - value=prompt.extract_graph_prompt, - label="Graph Extract Prompt Header", - language="markdown", - lines=15, - max_lines=29, ->>>>>>> 38dce0b (feat(llm): vector db finished) ) with gr.Row(): with gr.Column(): -<<<<<<< HEAD with gr.Tab("text") as tab_upload_text: input_text = gr.Textbox( value=prompt.doc_input_text, @@ -334,7 +291,9 @@ def create_vector_graph_block(): lines=15, max_lines=29, ) - out = gr.Code(label="Output Info", language="json", elem_classes="code-container-edit") + out = gr.Code( + label="Output Info", language="json", elem_classes="code-container-edit" + ) with gr.Row(): with gr.Accordion("Get RAG Info", open=False): @@ -346,31 +305,16 @@ def create_vector_graph_block(): vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") -======= - vector_index_btn0 = gr.Button("Get Vector Index Info", size="sm") - graph_index_btn0 = gr.Button("Get Graph Index Info", size="sm") - with gr.Accordion("Clear RAG Data", open=False): - with gr.Column(): - vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") - graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") - graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") - vector_import_bt = gr.Button("Import into Vector", variant="primary") - graph_extract_bt = gr.Button("Extract Graph Data (1)", variant="primary") - graph_loading_bt = gr.Button("Load into GraphDB (2)", interactive=True) - graph_index_rebuild_bt = gr.Button("Update Vid Embedding") ->>>>>>> 38dce0b (feat(llm): vector db finished) vector_import_bt = gr.Button("Import into Vector", variant="primary") graph_extract_bt = gr.Button("Extract Graph Data (1)", variant="primary") graph_loading_bt = gr.Button("Load into GraphDB (2)", interactive=True) graph_index_rebuild_bt = gr.Button("Update Vid Embedding") -<<<<<<< HEAD gr.Markdown("---") with gr.Accordion("Graph Schema Generator", open=False): gr.Markdown( - "Provide **query examples** and **few-shot examples**, " - "then click **Generate Schema** to automatically create graph schema." + "Provide **query examples** and **few-shot examples**, then click **Generate Schema** to automatically create graph schema." ) with gr.Row(): query_example = gr.Code( @@ -388,9 +332,10 @@ def create_vector_graph_block(): max_lines=15, ) build_schema_bt = gr.Button("Generate Schema", variant="primary") + _create_prompt_helper_block(demo, input_text, info_extract_template) - vector_index_btn0.click(get_vector_index_info, outputs=out).then( + vector_index_btn0.click(get_graph_index_info, outputs=out).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) @@ -420,24 +365,7 @@ def create_vector_graph_block(): store_prompt, inputs=[input_text, input_schema, info_extract_template], ) -======= - # origin_out = gr.Textbox(visible=False) - graph_extract_bt.click( - extract_graph, inputs=[input_file, input_text, input_schema, info_extract_template], outputs=[out] - ).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) - - graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( - update_vid_embedding - ).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) ->>>>>>> 38dce0b (feat(llm): vector db finished) - - # origin_out = gr.Textbox(visible=False) + graph_extract_bt.click( extract_graph, inputs=[input_file, input_text, input_schema, info_extract_template], @@ -447,25 +375,20 @@ def create_vector_graph_block(): inputs=[input_text, input_schema, info_extract_template], ) - graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( - update_vid_embedding - ).then( + graph_loading_bt.click( + import_graph_data, inputs=[out, input_schema], outputs=[out] + ).then(update_vid_embedding).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) - # TODO: we should store the examples after the user changed them. build_schema_bt.click( _build_schema_and_provide_feedback, inputs=[input_text, query_example, few_shot], outputs=[input_schema], ).then( store_prompt, - inputs=[ - input_text, - input_schema, - info_extract_template, - ], # TODO: Store the updated examples + inputs=[input_text, input_schema, info_extract_template], ) def on_tab_select(input_f, input_t, evt: gr.SelectData): @@ -507,7 +430,7 @@ async def timely_update_vid_embedding(interval_seconds: int = 3600): "pwd": huge_settings.graph_pwd, "graph_space": huge_settings.graph_space, } - if check_graph_db_connection(**config): # type:ignore + if check_graph_db_connection(**config): await asyncio.to_thread(update_vid_embedding) log.info("update_vid_embedding executed successfully") else: @@ -522,166 +445,4 @@ async def timely_update_vid_embedding(interval_seconds: int = 3600): # pylint: disable=W0718 except Exception as e: log.warning("Failed to execute update_vid_embedding: %s", e, exc_info=True) - await asyncio.sleep(interval_seconds) -<<<<<<< HEAD -======= - - -def create_vector_graph_block(): - # pylint: disable=no-member - # pylint: disable=C0301 - # pylint: disable=unexpected-keyword-arg - gr.Markdown( - """## Build Vector/Graph Index & Extract Knowledge Graph -- Docs: - - text: Build rag index from plain text - - file: Upload file(s) which should be TXT or .docx (Multiple files can be selected together) -- [Schema](https://hugegraph.apache.org/docs/clients/restful-api/schema/): (Accept **2 types**) - - User-defined Schema (JSON format, follow the [template](https://github.com/apache/incubator-hugegraph-ai/blob/aff3bbe25fa91c3414947a196131be812c20ef11/hugegraph-llm/src/hugegraph_llm/config/config_data.py#L125) - to modify it) - - Specify the name of the HugeGraph graph instance, it will automatically get the schema from it (like - **"hugegraph"**) -- Graph Extract Prompt Header: The user-defined prompt of graph extracting -- If already exist the graph data, you should click "**Rebuild vid Index**" to update the index -""" - ) - - with gr.Row(): - with gr.Column(): - with gr.Tab("text") as tab_upload_text: - input_text = gr.Textbox( - value=prompt.doc_input_text, label="Input Doc(s)", lines=20, show_copy_button=True - ) - with gr.Tab("file") as tab_upload_file: - input_file = gr.File( - value=None, - label="Docs (multi-files can be selected together)", - file_count="multiple", - ) - input_schema = gr.Code(value=prompt.graph_schema, label="Graph Schema", language="json", lines=15, max_lines=29) - info_extract_template = gr.Code( - value=prompt.extract_graph_prompt, - label="Graph Extract Prompt Header", - language="markdown", - lines=15, - max_lines=29, - ) - - out = gr.Code(label="Output Info", language="json", elem_classes="code-container-edit") - - with gr.Row(): - with gr.Accordion("Get RAG Info", open=False): - with gr.Column(): - vector_index_btn0 = gr.Button("Get Vector Index Info", size="sm") - graph_index_btn0 = gr.Button("Get Graph Index Info", size="sm") - with gr.Accordion("Clear RAG Data", open=False): - with gr.Column(): - vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") - graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") - graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") - vector_import_bt = gr.Button("Import into Vector", variant="primary") - graph_extract_bt = gr.Button("Extract Graph Data (1)", variant="primary") - graph_loading_bt = gr.Button("Load into GraphDB (2)", interactive=True) - graph_index_rebuild_bt = gr.Button("Update Vid Embedding") - - gr.Markdown("---") - with gr.Accordion("Graph Schema Generator", open=False): - gr.Markdown( - "Provide **query examples** and **few-shot examples**, " - "then click **Generate Schema** to automatically create graph schema." - ) - with gr.Row(): - query_example = gr.Code( - value=load_query_examples(), - label="Query Examples", - language="json", - lines=10, - max_lines=15, - ) - few_shot = gr.Code( - value=load_schema_fewshot_examples(), - label="Few-shot Example", - language="json", - lines=10, - max_lines=15, - ) - build_schema_bt = gr.Button("Generate Schema", variant="primary") - - # 事件绑定 - vector_index_btn0.click(get_vector_index_info, outputs=out).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) - vector_index_btn1.click(clean_vector_index).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) - vector_import_bt.click(build_vector_index, inputs=[input_file, input_text], outputs=out).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) - graph_index_btn0.click(get_graph_index_info, outputs=out).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) - graph_index_btn1.click(clean_all_graph_index).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) - graph_data_btn0.click(clean_all_graph_data).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) - graph_index_rebuild_bt.click(update_vid_embedding, outputs=out).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) - - graph_extract_bt.click( - extract_graph, inputs=[input_file, input_text, input_schema, info_extract_template], outputs=[out] - ).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) - - graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( - update_vid_embedding - ).then( - store_prompt, - inputs=[input_text, input_schema, info_extract_template], - ) - - build_schema_bt.click( - lambda it, qe, fs: extract_graph([], it, prompt.graph_schema, prompt.extract_graph_prompt), - inputs=[input_text, query_example, few_shot], - outputs=[input_schema], - ).then( - store_prompt, - inputs=[ - input_text, - input_schema, - info_extract_template, - ], - ) - - def on_tab_select(input_f, input_t, evt: gr.SelectData): - print(f"You selected {evt.value} at {evt.index} from {evt.target}") - if evt.value == "file": - return input_f, "" - if evt.value == "text": - return [], input_t - return [], "" - - tab_upload_file.select( - fn=on_tab_select, - inputs=[input_file, input_text], - outputs=[input_file, input_text], - ) - tab_upload_text.select( - fn=on_tab_select, - inputs=[input_file, input_text], - outputs=[input_file, input_text], - ) - - return input_text, input_schema, info_extract_template ->>>>>>> 87ee5d3 (style: format code with black line-length 120) + await asyncio.sleep(interval_seconds) \ No newline at end of file diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py index ed3f28a2a..64967d2df 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py @@ -15,76 +15,33 @@ # specific language governing permissions and limitations # under the License. - -<<<<<<< HEAD import asyncio -import os -from typing import Dict, Any -======= from typing import Any, Dict ->>>>>>> 38dce0b (feat(llm): vector db finished) - -<<<<<<< HEAD -from hugegraph_llm.config import huge_settings, resource_path, llm_settings -from hugegraph_llm.indices.vector_index import VectorIndex -======= -from tqdm import tqdm -<<<<<<< HEAD -from hugegraph_llm.config import huge_settings, resource_path -from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) -======= from hugegraph_llm.config import huge_settings from hugegraph_llm.indices.vector_index.base import VectorStoreBase ->>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import ( - get_embeddings_parallel, - get_filename_prefix, - get_index_folder_name, -) +from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel from hugegraph_llm.utils.log import log class BuildVectorIndex: def __init__(self, embedding: BaseEmbedding, vector_index: type[VectorStoreBase]): self.embedding = embedding -<<<<<<< HEAD -<<<<<<< HEAD - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(self.embedding, "model_name", None) - ) - self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) -======= - self.index_dir = str(os.path.join(resource_path, huge_settings.graph_name, "chunks")) - self.vector_index = FaissVectorIndex.from_name(self.index_dir) ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) -======= self.vector_index = vector_index.from_name( embedding.get_embedding_dim(), huge_settings.graph_name, "chunks", ) ->>>>>>> 38dce0b (feat(llm): vector db finished) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: if "chunks" not in context: raise ValueError("chunks not found in context.") chunks = context["chunks"] - chunks_embedding = [] log.debug("Building vector index for %s chunks...", len(context["chunks"])) - # TODO: use async_get_texts_embedding instead of single sync method - chunks_embedding = asyncio.run(get_embeddings_parallel(self.embedding, chunks)) + # Use async parallel embedding to speed up + chunks_embedding = asyncio.run(get_embeddings_parallel(self.embedding, chunks)) # type: ignore if len(chunks_embedding) > 0: self.vector_index.add(chunks_embedding, chunks) -<<<<<<< HEAD - self.vector_index.to_index_file(self.index_dir, self.filename_prefix) -======= self.vector_index.save_index_by_name(huge_settings.graph_name, "chunks") ->>>>>>> 38dce0b (feat(llm): vector db finished) return context diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py index 0443b15c8..0b396b77d 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py @@ -34,28 +34,18 @@ from hugegraph_llm.operators.llm_op.disambiguate_data import DisambiguateData from hugegraph_llm.operators.llm_op.info_extract import InfoExtract from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract -<<<<<<< HEAD from hugegraph_llm.operators.llm_op.schema_build import SchemaBuilder -from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm -from pyhugegraph.client import PyHugeClient -======= from hugegraph_llm.utils.decorators import log_operator_time, log_time, record_rpm ->>>>>>> 38dce0b (feat(llm): vector db finished) class KgBuilder: -<<<<<<< HEAD def __init__( self, llm: BaseLLM, embedding: Optional[BaseEmbedding] = None, graph: Optional[PyHugeClient] = None, ): - self.operators = [] -======= - def __init__(self, llm: BaseLLM, embedding: Optional[BaseEmbedding] = None, graph: Optional[PyHugeClient] = None): self.operators: List[Any] = [] ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) self.llm = llm self.embedding = embedding self.graph = graph @@ -119,8 +109,8 @@ def print_result(self): self.operators.append(PrintResult()) return self - def build_schema(self): - self.operators.append(SchemaBuilder(self.llm)) + def build_schema(self, schema_prompt: Optional[str] = None): + self.operators.append(SchemaBuilder(self.llm, schema_prompt)) return self @log_time("total time") From 17c72bc082fb22845bf14e37e978ebaa2e91995a Mon Sep 17 00:00:00 2001 From: lingxiao Date: Wed, 3 Sep 2025 15:56:59 +0800 Subject: [PATCH 28/71] fix black --- .../hugegraph_llm/api/models/rag_requests.py | 24 +++-- .../src/hugegraph_llm/api/rag_api.py | 23 +++-- .../src/hugegraph_llm/config/llm_config.py | 16 +++- .../config/models/base_config.py | 12 ++- .../config/models/base_prompt_config.py | 12 ++- .../demo/rag_demo/admin_block.py | 4 +- .../src/hugegraph_llm/demo/rag_demo/app.py | 4 +- .../demo/rag_demo/configs_block.py | 87 ++++++++++++++----- .../demo/rag_demo/other_block.py | 12 ++- .../hugegraph_llm/demo/rag_demo/rag_block.py | 40 ++++++--- .../demo/rag_demo/text2gremlin_block.py | 34 ++++++-- .../demo/rag_demo/vector_graph_block.py | 22 +++-- .../vector_index/milvus_vector_store.py | 8 +- .../vector_index/qdrant_vector_store.py | 8 +- .../hugegraph_llm/middleware/middleware.py | 4 +- .../models/embeddings/init_embedding.py | 8 +- .../hugegraph_llm/models/embeddings/openai.py | 4 +- .../src/hugegraph_llm/models/llms/init_llm.py | 6 +- .../src/hugegraph_llm/models/llms/ollama.py | 8 +- .../src/hugegraph_llm/models/llms/openai.py | 12 ++- .../hugegraph_llm/models/rerankers/cohere.py | 4 +- .../models/rerankers/siliconflow.py | 4 +- .../operators/common_op/check_schema.py | 40 ++++++--- .../operators/common_op/merge_dedup_rerank.py | 10 ++- .../operators/common_op/nltk_helper.py | 4 +- .../operators/document_op/word_extract.py | 6 +- .../hugegraph_llm/operators/graph_rag_task.py | 4 +- .../operators/gremlin_generate_task.py | 16 +++- .../hugegraph_op/commit_to_hugegraph.py | 62 +++++++++---- .../hugegraph_op/fetch_graph_data.py | 4 +- .../operators/hugegraph_op/graph_rag_query.py | 44 +++++++--- .../operators/hugegraph_op/schema_manager.py | 4 +- .../index_op/build_gremlin_example_index.py | 8 +- .../index_op/build_semantic_index.py | 20 +++-- .../index_op/gremlin_example_index_query.py | 14 +-- .../operators/index_op/semantic_id_query.py | 8 +- .../operators/kg_construction_task.py | 4 +- .../operators/llm_op/answer_synthesize.py | 60 +++++++++---- .../operators/llm_op/gremlin_generate.py | 16 +++- .../operators/llm_op/info_extract.py | 8 +- .../operators/llm_op/keyword_extract.py | 4 +- .../operators/llm_op/prompt_generate.py | 12 ++- .../llm_op/property_graph_extract.py | 18 ++-- .../operators/llm_op/schema_build.py | 8 +- .../llm_op/unstructured_data_utils.py | 8 +- .../src/hugegraph_llm/utils/decorators.py | 4 +- .../hugegraph_llm/utils/embedding_utils.py | 12 ++- .../hugegraph_llm/utils/graph_index_utils.py | 24 +++-- .../hugegraph_llm/utils/hugegraph_utils.py | 40 ++++++--- .../hugegraph_llm/utils/vector_index_utils.py | 14 ++- .../embeddings/test_ollama_embedding.py | 4 +- .../operators/common_op/test_check_schema.py | 4 +- .../llm_op/test_disambiguate_data.py | 4 +- 53 files changed, 626 insertions(+), 218 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py index f46aea02c..9ebc94eb1 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py +++ b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py @@ -34,12 +34,18 @@ class GraphConfigRequest(BaseModel): class RAGRequest(BaseModel): query: str = Query(..., description="Query you want to ask") raw_answer: bool = Query(False, description="Use LLM to generate answer directly") - vector_only: bool = Query(False, description="Use LLM to generate answer with vector") - graph_only: bool = Query(True, description="Use LLM to generate answer with graph RAG only") + vector_only: bool = Query( + False, description="Use LLM to generate answer with vector" + ) + graph_only: bool = Query( + True, description="Use LLM to generate answer with graph RAG only" + ) graph_vector_answer: bool = Query( False, description="Use LLM to generate answer with vector & GraphRAG" ) - graph_ratio: float = Query(0.5, description="The ratio of GraphRAG ans & vector ans") + graph_ratio: float = Query( + 0.5, description="The ratio of GraphRAG ans & vector ans" + ) rerank_method: Literal["bleu", "reranker"] = Query( "bleu", description="Method to rerank the results." ) @@ -53,7 +59,9 @@ class RAGRequest(BaseModel): max_graph_items: int = Query( 30, description="Maximum number of items for GQL queries in graph." ) - topk_return_results: int = Query(20, description="Number of sorted results to return finally.") + topk_return_results: int = Query( + 20, description="Number of sorted results to return finally." + ) vector_dis_threshold: float = Query( 0.9, description="Threshold for vector similarity\ @@ -90,7 +98,9 @@ class GraphRAGRequest(BaseModel): max_graph_items: int = Query( 30, description="Maximum number of items for GQL queries in graph." ) - topk_return_results: int = Query(20, description="Number of sorted results to return finally.") + topk_return_results: int = Query( + 20, description="Number of sorted results to return finally." + ) vector_dis_threshold: float = Query( 0.9, description="Threshold for vector similarity \ @@ -105,7 +115,9 @@ class GraphRAGRequest(BaseModel): client_config: Optional[GraphConfigRequest] = Query( None, description="hugegraph server config." ) - get_vertex_only: bool = Query(False, description="return only keywords & vertex (early stop).") + get_vertex_only: bool = Query( + False, description="return only keywords & vertex (early stop)." + ) gremlin_tmpl_num: int = Query( 1, diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py index 5c9295efa..1d5b451b1 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py @@ -67,7 +67,8 @@ def rag_answer_api(req: RAGRequest): # Keep prompt params in the end custom_related_information=req.custom_priority_info, answer_prompt=req.answer_prompt or prompt.answer_prompt, - keywords_extract_prompt=req.keywords_extract_prompt or prompt.keywords_extract_prompt, + keywords_extract_prompt=req.keywords_extract_prompt + or prompt.keywords_extract_prompt, gremlin_prompt=req.gremlin_prompt or prompt.gremlin_generate_prompt, ) # TODO: we need more info in the response for users to understand the query logic @@ -135,7 +136,9 @@ def graph_rag_recall_api(req: GraphRAGRequest): except TypeError as e: log.error("TypeError in graph_rag_recall_api: %s", e) - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) + ) from e except Exception as e: log.error("Unexpected error occurred: %s", e) raise HTTPException( @@ -146,7 +149,9 @@ def graph_rag_recall_api(req: GraphRAGRequest): @router.post("/config/graph", status_code=status.HTTP_201_CREATED) def graph_config_api(req: GraphConfigRequest): # Accept status code - res = apply_graph_conf(req.url, req.name, req.user, req.pwd, req.gs, origin_call="http") + res = apply_graph_conf( + req.url, req.name, req.user, req.pwd, req.gs, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) # TODO: restructure the implement of llm to three types, like "/config/chat_llm" @@ -159,7 +164,9 @@ def llm_config_api(req: LLMConfigRequest): req.api_key, req.api_base, req.language_model, req.max_tokens, origin_call="http" ) else: - res = apply_llm_conf(req.host, req.port, req.language_model, None, origin_call="http") + res = apply_llm_conf( + req.host, req.port, req.language_model, None, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/embedding", status_code=status.HTTP_201_CREATED) @@ -171,7 +178,9 @@ def embedding_config_api(req: LLMConfigRequest): req.api_key, req.api_base, req.language_model, origin_call="http" ) else: - res = apply_embedding_conf(req.host, req.port, req.language_model, origin_call="http") + res = apply_embedding_conf( + req.host, req.port, req.language_model, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/rerank", status_code=status.HTTP_201_CREATED) @@ -183,7 +192,9 @@ def rerank_config_api(req: RerankerConfigRequest): req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http" ) elif req.reranker_type == "siliconflow": - res = apply_reranker_conf(req.api_key, req.reranker_model, None, origin_call="http") + res = apply_reranker_conf( + req.api_key, req.reranker_model, None, origin_call="http" + ) else: res = status.HTTP_501_NOT_IMPLEMENTED return generate_response(RAGResponse(status_code=res, message="Missing Value")) diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index af22b71b0..8b4f274ee 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -32,13 +32,19 @@ class LLMConfig(BaseConfig): embedding_type: Optional[Literal["openai", "litellm", "ollama/local"]] = "openai" reranker_type: Optional[Literal["cohere", "siliconflow"]] = None # 1. OpenAI settings - openai_chat_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_chat_api_base: str = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_chat_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_chat_language_model: str = "gpt-4.1-mini" - openai_extract_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_extract_api_base: str = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_extract_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_extract_language_model: str = "gpt-4.1-mini" - openai_text2gql_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_text2gql_api_base: str = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_text2gql_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_text2gql_language_model: str = "gpt-4.1-mini" openai_embedding_api_base: str = os.environ.get( @@ -51,7 +57,9 @@ class LLMConfig(BaseConfig): openai_extract_tokens: int = 256 openai_text2gql_tokens: int = 4096 # 2. Rerank settings - cohere_base_url: str = os.environ.get("CO_API_URL", "https://api.cohere.com/v1/rerank") + cohere_base_url: str = os.environ.get( + "CO_API_URL", "https://api.cohere.com/v1/rerank" + ) reranker_api_key: Optional[str] = None reranker_model: Optional[str] = None # 3. Ollama settings diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py index 5fec3a778..4ec9256c5 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py @@ -24,7 +24,9 @@ from hugegraph_llm.utils.log import log dir_name = os.path.dirname -env_path = os.path.join(os.getcwd(), ".env") # Load .env from the current working directory +env_path = os.path.join( + os.getcwd(), ".env" +) # Load .env from the current working directory class BaseConfig(BaseSettings): @@ -88,7 +90,9 @@ def check_env(self): # Step 2: Add missing config items to .env self._sync_object_to_env(env_config, config_dict) except Exception as e: - log.error("An error occurred when checking the .env variable file: %s", str(e)) + log.error( + "An error occurred when checking the .env variable file: %s", str(e) + ) raise def _sync_env_to_object(self, env_config, config_dict): @@ -139,7 +143,9 @@ def __init__(self, **data): # Synchronize configurations between the object and .env file self.check_env() - log.info("The %s file was loaded. Class: %s", env_path, self.__class__.__name__) + log.info( + "The %s file was loaded. Class: %s", env_path, self.__class__.__name__ + ) except Exception as e: log.error("An error occurred when initializing the configuration object: %s", str(e)) raise diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py index 2369d01a6..b15bad0a6 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py @@ -57,7 +57,9 @@ def ensure_yaml_file_exists(self): current_dir = Path.cwd().resolve() project_root = get_project_root() if current_dir == project_root: - log.info("Current working directory is the project root, proceeding to run the app.") + log.info( + "Current working directory is the project root, proceeding to run the app." + ) else: error_msg = ( f"Current working directory is not the project root. " @@ -122,7 +124,9 @@ def to_literal(val): "gremlin_generate_prompt": to_literal(self.gremlin_generate_prompt), "doc_input_text": to_literal(self.doc_input_text), "_language_generated": str(self.llm_settings.language).lower().strip(), - "generate_extract_prompt_template": to_literal(self.generate_extract_prompt_template), + "generate_extract_prompt_template": to_literal( + self.generate_extract_prompt_template + ), } with open(yaml_file_path, "w", encoding="utf-8") as file: yaml.dump(data, file, allow_unicode=True, sort_keys=False, default_flow_style=False) @@ -150,7 +154,9 @@ def generate_yaml_file(self): self.keywords_extract_prompt = self.keywords_extract_prompt_EN self.doc_input_text = self.doc_input_text_EN self.save_to_yaml() - log.info("Prompt file '%s' has been generated with default values.", yaml_file_path) + log.info( + "Prompt file '%s' has been generated with default values.", yaml_file_path + ) def update_yaml_file(self): self.save_to_yaml() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 1a157d387..d3beebcbb 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -143,7 +143,9 @@ def create_admin_block(): with gr.Row(): with gr.Column(): # Button to clear LLM Server log, initially hidden - clear_llm_server_button = gr.Button("Clear LLM Server Log", visible=False) + clear_llm_server_button = gr.Button( + "Clear LLM Server Log", visible=False + ) with gr.Column(): # Button to refresh LLM Server log manually refresh_llm_server_button = gr.Button( 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 b0979763c..a78e62361 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -166,7 +166,9 @@ def create_app(): prompt.update_yaml_file() assert admin_settings.enable_login auth_enabled = admin_settings.enable_login.lower() == "true" - log.info("(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled") + log.info( + "(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled" + ) api_auth = APIRouter(dependencies=[Depends(authenticate)] if auth_enabled else []) hugegraph_llm = init_rag_ui() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 112892e7c..721cf272c 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -71,7 +71,9 @@ def test_api_connection( log.debug("Request URL: %s", url) try: if method.upper() == "GET": - resp = requests.get(url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth) + resp = requests.get( + url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth + ) elif method.upper() == "POST": resp = requests.post( url, @@ -106,7 +108,9 @@ def test_api_connection( return resp.status_code -def config_qianfan_model(arg1, arg2, arg3=None, settings_prefix=None, origin_call=None) -> int: +def config_qianfan_model( + arg1, arg2, arg3=None, settings_prefix=None, origin_call=None +) -> int: setattr(llm_settings, f"qianfan_{settings_prefix}_api_key", arg1) setattr(llm_settings, f"qianfan_{settings_prefix}_secret_key", arg2) if arg3: @@ -142,7 +146,9 @@ def apply_embedding_config(arg1, arg2, arg3, origin_call=None) -> int: llm_settings.ollama_embedding_host = arg1 llm_settings.ollama_embedding_port = int(arg2) llm_settings.ollama_embedding_model = arg3 - status_code = test_api_connection(f"http://{arg1}:{arg2}", origin_call=origin_call) + status_code = test_api_connection( + f"http://{arg1}:{arg2}", origin_call=origin_call + ) elif embedding_option == "litellm": llm_settings.litellm_embedding_api_key = arg1 llm_settings.litellm_embedding_api_base = arg2 @@ -233,7 +239,8 @@ def apply_llm_config( setattr(llm_settings, f"openai_{current_llm_config}_tokens", int(max_tokens)) test_url = ( - getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" + getattr(llm_settings, f"openai_{current_llm_config}_api_base") + + "/chat/completions" ) data = { "model": model_name, @@ -247,7 +254,9 @@ def apply_llm_config( elif llm_option == "ollama/local": setattr(llm_settings, f"ollama_{current_llm_config}_host", api_key_or_host) - setattr(llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port)) + setattr( + llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port) + ) setattr(llm_settings, f"ollama_{current_llm_config}_language_model", model_name) status_code = test_api_connection( f"http://{api_key_or_host}:{api_base_or_port}", origin_call=origin_call @@ -255,8 +264,12 @@ def apply_llm_config( elif llm_option == "litellm": setattr(llm_settings, f"litellm_{current_llm_config}_api_key", api_key_or_host) - setattr(llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port) - setattr(llm_settings, f"litellm_{current_llm_config}_language_model", model_name) + setattr( + llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port + ) + setattr( + llm_settings, f"litellm_{current_llm_config}_language_model", model_name + ) setattr(llm_settings, f"litellm_{current_llm_config}_tokens", int(max_tokens)) status_code = test_litellm_chat( @@ -383,9 +396,13 @@ def chat_llm_settings(llm_type): ), ] else: - llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] + llm_config_input = [ + gr.Textbox(value="", visible=False) for _ in range(4) + ] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click(apply_llm_config_with_chat_op, inputs=llm_config_input) + llm_config_button.click( + apply_llm_config_with_chat_op, inputs=llm_config_input + ) # Determine whether there are Settings in the.env file env_path = os.path.join( os.getcwd(), ".env" @@ -425,7 +442,9 @@ def extract_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr(llm_settings, "openai_extract_language_model"), + value=getattr( + llm_settings, "openai_extract_language_model" + ), label="model_name", ), gr.Textbox( @@ -444,7 +463,9 @@ def extract_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr(llm_settings, "ollama_extract_language_model"), + value=getattr( + llm_settings, "ollama_extract_language_model" + ), label="model_name", ), gr.Textbox(value="", visible=False), @@ -462,7 +483,9 @@ def extract_llm_settings(llm_type): info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr(llm_settings, "litellm_extract_language_model"), + value=getattr( + llm_settings, "litellm_extract_language_model" + ), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -472,9 +495,13 @@ def extract_llm_settings(llm_type): ), ] else: - llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] + llm_config_input = [ + gr.Textbox(value="", visible=False) for _ in range(4) + ] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) + llm_config_button.click( + apply_llm_config_with_extract_op, inputs=llm_config_input + ) with gr.Tab(label="text2gql"): text2gql_llm_dropdown = gr.Dropdown( @@ -499,7 +526,9 @@ def text2gql_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr(llm_settings, "openai_text2gql_language_model"), + value=getattr( + llm_settings, "openai_text2gql_language_model" + ), label="model_name", ), gr.Textbox( @@ -518,7 +547,9 @@ def text2gql_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr(llm_settings, "ollama_text2gql_language_model"), + value=getattr( + llm_settings, "ollama_text2gql_language_model" + ), label="model_name", ), gr.Textbox(value="", visible=False), @@ -536,7 +567,9 @@ def text2gql_llm_settings(llm_type): info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr(llm_settings, "litellm_text2gql_language_model"), + value=getattr( + llm_settings, "litellm_text2gql_language_model" + ), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -546,9 +579,13 @@ def text2gql_llm_settings(llm_type): ), ] else: - llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] + llm_config_input = [ + gr.Textbox(value="", visible=False) for _ in range(4) + ] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click(apply_llm_config_with_text2gql_op, inputs=llm_config_input) + llm_config_button.click( + apply_llm_config_with_text2gql_op, inputs=llm_config_input + ) with gr.Accordion("3. Set up the Embedding.", open=False): embedding_dropdown = gr.Dropdown( @@ -635,7 +672,9 @@ def embedding_settings(embedding_type): @gr.render(inputs=[reranker_dropdown]) def reranker_settings(reranker_type): - llm_settings.reranker_type = reranker_type if reranker_type != "None" else None + llm_settings.reranker_type = ( + reranker_type if reranker_type != "None" else None + ) if reranker_type == "cohere": with gr.Row(): reranker_config_input = [ @@ -644,8 +683,12 @@ def reranker_settings(reranker_type): label="api_key", type="password", ), - gr.Textbox(value=lambda: llm_settings.reranker_model, label="model"), - gr.Textbox(value=lambda: llm_settings.cohere_base_url, label="base_url"), + gr.Textbox( + value=lambda: llm_settings.reranker_model, label="model" + ), + gr.Textbox( + value=lambda: llm_settings.cohere_base_url, label="base_url" + ), ] elif reranker_type == "siliconflow": with gr.Row(): diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py index 82650b907..3f8089b6d 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py @@ -34,9 +34,13 @@ def create_other_block(): inp = gr.Textbox( value="g.V().limit(10)", label="Gremlin query", show_copy_button=True, lines=8 ) - out = gr.Code(label="Output", language="json", elem_classes="code-container-show") + out = gr.Code( + label="Output", language="json", elem_classes="code-container-show" + ) btn = gr.Button("Run Gremlin query") - btn.click(fn=run_gremlin_query, inputs=[inp], outputs=out) # pylint: disable=no-member + btn.click( + fn=run_gremlin_query, inputs=[inp], outputs=out + ) # pylint: disable=no-member gr.Markdown("---") with gr.Row(): @@ -51,7 +55,9 @@ def create_other_block(): inp = [] out = gr.Textbox(label="Init Graph Demo Result", show_copy_button=True) btn = gr.Button("(BETA) Init HugeGraph test data (🚧)") - btn.click(fn=init_hg_test_data, inputs=inp, outputs=out) # pylint: disable=no-member + btn.click( + fn=init_hg_test_data, inputs=inp, outputs=out + ) # pylint: disable=no-member @asynccontextmanager diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 2a817ada4..31a2d3b09 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -114,7 +114,9 @@ def rag_answer( max_graph_items=max_graph_items, ) if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") + gr.Warning( + "Online reranker fails, automatically switches to local bleu rerank." + ) return ( context.get("raw_answer", ""), context.get("vector_only_answer", ""), @@ -222,7 +224,9 @@ async def rag_answer_streaming( graph_search=graph_search, ) if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") + gr.Warning( + "Online reranker fails, automatically switches to local bleu rerank." + ) answer_synthesize = AnswerSynthesize( raw_answer=raw_answer, vector_only_answer=vector_only_answer, @@ -232,7 +236,9 @@ async def rag_answer_streaming( ) async for context in answer_synthesize.run_streaming(context): if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") + gr.Warning( + "Online reranker fails, automatically switches to local bleu rerank." + ) yield ( context.get("raw_answer", ""), context.get("vector_only_answer", ""), @@ -302,7 +308,9 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): - raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") + raw_radio = gr.Radio( + choices=[True, False], value=False, label="Basic LLM Answer" + ) vector_only_radio = gr.Radio( choices=[True, False], value=False, label="Vector-only Answer" ) @@ -387,7 +395,9 @@ def toggle_slider(enable): # FIXME: "demo" might conflict with the graph name, it should be modified. answers_path = os.path.join(resource_path, "demo", "questions_answers.xlsx") questions_path = os.path.join(resource_path, "demo", "questions.xlsx") - questions_template_path = os.path.join(resource_path, "demo", "questions_template.xlsx") + questions_template_path = os.path.join( + resource_path, "demo", "questions_template.xlsx" + ) def read_file_to_excel(file: Any, line_count: Optional[int] = None): df = None @@ -467,12 +477,18 @@ def several_rag_answer( file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" ) with gr.Column(): - test_template_file = os.path.join(resource_path, "demo", "questions_template.xlsx") + test_template_file = os.path.join( + resource_path, "demo", "questions_template.xlsx" + ) gr.File(value=test_template_file, label="Download Template File") - answer_max_line_count = gr.Number(1, label="Max Lines To Show", minimum=1, maximum=40) + answer_max_line_count = gr.Number( + 1, label="Max Lines To Show", minimum=1, maximum=40 + ) answers_btn = gr.Button("Generate Answer (Batch)", variant="primary") # TODO: Set individual progress bars for dataframe - qa_dataframe = gr.DataFrame(label="Questions & Answers (Preview)", headers=tests_df_headers) + qa_dataframe = gr.DataFrame( + label="Questions & Answers (Preview)", headers=tests_df_headers + ) answers_btn.click( several_rag_answer, inputs=[ @@ -490,8 +506,12 @@ def several_rag_answer( ], outputs=[qa_dataframe, gr.File(label="Download Answered File", min_width=40)], ) - questions_file.change(read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count]) - answer_max_line_count.change(change_showing_excel, answer_max_line_count, qa_dataframe) + questions_file.change( + read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count] + ) + answer_max_line_count.change( + change_showing_excel, answer_max_line_count, qa_dataframe + ) return ( inp, answer_prompt_input, diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 8fbb01c25..af55cd86f 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -96,8 +96,12 @@ def build_example_vector_index(temp_file) -> dict: timestamp = datetime.now().strftime("%Y%m%d%H%M%S") _, file_name = os.path.split(f"{name}_{timestamp}{ext}") log.info("Copying file to: %s", file_name) - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - target_file = os.path.join(resource_path, folder_name, "gremlin_examples", file_name) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + target_file = os.path.join( + resource_path, folder_name, "gremlin_examples", file_name + ) try: import shutil @@ -197,9 +201,13 @@ def gremlin_generate( processed_schema, short_schema = _process_schema(schema, generator, sm) if processed_schema is None and short_schema is None: - return GremlinResult.error("Invalid JSON schema, please check the format carefully.") + return GremlinResult.error( + "Invalid JSON schema, please check the format carefully." + ) - updated_schema = sm.simple_schema(processed_schema) if short_schema else processed_schema + updated_schema = ( + sm.simple_schema(processed_schema) if short_schema else processed_schema + ) store_schema(str(updated_schema), inp, gremlin_prompt) output_types = _configure_output_types(requested_outputs) @@ -231,7 +239,11 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} + new_vertex = { + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex + } mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) @@ -282,7 +294,9 @@ def create_text2gremlin_block() -> Tuple: with gr.Row(): btn = gr.Button("Build Example Vector Index", variant="primary") - btn.click(build_example_vector_index, inputs=[file], outputs=[out]) # pylint: disable=no-member + btn.click( + build_example_vector_index, inputs=[file], outputs=[out] + ) # pylint: disable=no-member gr.Markdown("## Nature Language To Gremlin") with gr.Row(): @@ -297,8 +311,12 @@ def create_text2gremlin_block() -> Tuple: language="javascript", elem_classes="code-container-show", ) - initialized_out = gr.Textbox(label="Gremlin With Template", show_copy_button=True) - raw_out = gr.Textbox(label="Gremlin Without Template", show_copy_button=True) + initialized_out = gr.Textbox( + label="Gremlin With Template", show_copy_button=True + ) + raw_out = gr.Textbox( + label="Gremlin Without Template", show_copy_button=True + ) tmpl_exec_out = gr.Code( label="Query With Template Output", language="json", diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 1e8124b81..c5f47d0b5 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -103,9 +103,11 @@ def load_query_examples(): language = getattr( prompt, "language", + ( getattr(prompt.llm_settings, "language", "EN") if hasattr(prompt, "llm_settings") - else "EN", + else "EN" + ), ) if language.upper() == "CN": examples_path = os.path.join( @@ -182,9 +184,11 @@ def _create_prompt_helper_block(demo, input_text, info_extract_template): few_shot_dropdown = gr.Dropdown( choices=example_names, label="Select a Few-shot example as a reference", - value=example_names[0] - if example_names and example_names[0] != "No available examples" - else None, + value=( + example_names[0] + if example_names and example_names[0] != "No available examples" + else None + ), ) with gr.Accordion("View example details", open=False): example_desc_preview = gr.Markdown(label="Example description") @@ -302,8 +306,12 @@ def create_vector_graph_block(): graph_index_btn0 = gr.Button("Get Graph Index Info", size="sm") with gr.Accordion("Clear RAG Data", open=False): with gr.Column(): - vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") - graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") + vector_index_btn1 = gr.Button( + "Clear Chunks Vector Index", size="sm" + ) + graph_index_btn1 = gr.Button( + "Clear Graph Vid Vector Index", size="sm" + ) graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") vector_import_bt = gr.Button("Import into Vector", variant="primary") @@ -445,4 +453,4 @@ async def timely_update_vid_embedding(interval_seconds: int = 3600): # pylint: disable=W0718 except Exception as e: log.warning("Failed to execute update_vid_embedding: %s", e, exc_info=True) - await asyncio.sleep(interval_seconds) \ No newline at end of file + await asyncio.sleep(interval_seconds) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py index b168b8c39..fae5ca860 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -73,11 +73,15 @@ def __init__( def _create_collection(self): """Create a new collection in Milvus.""" - id_field = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True) + id_field = FieldSchema( + name="id", dtype=DataType.INT64, is_primary=True, auto_id=True + ) vector_field = FieldSchema( name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embed_dim ) - property_field = FieldSchema(name="property", dtype=DataType.VARCHAR, max_length=65535) + property_field = FieldSchema( + name="property", dtype=DataType.VARCHAR, max_length=65535 + ) original_id_field = FieldSchema(name="original_id", dtype=DataType.INT64) schema = CollectionSchema( diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index ca7761ecd..14b97fbff 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -28,7 +28,9 @@ class QdrantVectorIndex(VectorStoreBase): - def __init__(self, name: str, host: str, port: int, api_key=None, embed_dim: int = 1024): + def __init__( + self, name: str, host: str, port: int, api_key=None, embed_dim: int = 1024 + ): self.embed_dim = embed_dim self.host = host self.port = port @@ -115,7 +117,9 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: return remove_num - def search(self, query_vector: List[float], top_k: int = 5, dis_threshold: float = 0.9): + def search( + self, query_vector: List[float], top_k: int = 5, dis_threshold: float = 0.9 + ): search_result = self.client.search( collection_name=self.name, query_vector=query_vector, limit=top_k ) diff --git a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py index f13e11e6f..5d98ebdf8 100644 --- a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py +++ b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py @@ -41,7 +41,9 @@ async def dispatch(self, request: Request, call_next): unit = "s" response.headers["X-Process-Time"] = f"{process_time:.2f} {unit}" - log.info("Request process time: %.2f ms, code=%d", process_time, response.status_code) + log.info( + "Request process time: %.2f ms, code=%d", process_time, response.status_code + ) log.info( "%s - Args: %s, IP: %s, URL: %s", request.method, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index d96840911..8e6af2774 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -65,7 +65,9 @@ def __init__(self): def get_embedding(self): if self.embedding_type == "openai": - assert llm_settings.openai_embedding_model_dim, "openai_embedding_model_dim is need" + assert ( + llm_settings.openai_embedding_model_dim + ), "openai_embedding_model_dim is need" return OpenAIEmbedding( embedding_dimension=llm_settings.openai_embedding_model_dim, model_name=llm_settings.openai_embedding_model, @@ -73,7 +75,9 @@ def get_embedding(self): api_base=llm_settings.openai_embedding_api_base, ) if self.embedding_type == "ollama/local": - assert llm_settings.ollama_embedding_model_dim, "ollama_embedding_model_dim is need" + assert ( + llm_settings.ollama_embedding_model_dim + ), "ollama_embedding_model_dim is need" return OllamaEmbedding( <<<<<<< HEAD model_name=llm_settings.ollama_embedding_model, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index 6e30cca71..15d928286 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -91,5 +91,7 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] A list of embedding vectors, where each vector is a list of floats. The order of embeddings should match the order of input texts. """ - response = await self.aclient.embeddings.create(input=texts, model=self.model_name) + response = await self.aclient.embeddings.create( + input=texts, model=self.model_name + ) return [data.embedding for data in response.data] 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 9121fca09..7e1eaab68 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py @@ -173,4 +173,8 @@ def get_text2gql_llm(self): if __name__ == "__main__": client = LLMs().get_chat_llm() print(client.generate(prompt="What is the capital of China?")) - print(client.generate(messages=[{"role": "user", "content": "What is the capital of China?"}])) + print( + client.generate( + messages=[{"role": "user", "content": "What is the capital of China?"}] + ) + ) diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py index 199384b12..c15c5440e 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py @@ -29,7 +29,9 @@ class OllamaClient(BaseLLM): """LLM wrapper should take in a prompt and return a string.""" - def __init__(self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwargs): + def __init__( + self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwargs + ): self.model = model self.client = ollama.Client(host=f"http://{host}:{port}", **kwargs) self.async_client = ollama.AsyncClient(host=f"http://{host}:{port}", **kwargs) @@ -99,7 +101,9 @@ def generate_streaming( for chunk in self.client.chat(model=self.model, messages=messages, stream=True): if not chunk["message"]: - log.debug("Received empty chunk['message'] in streaming chunk: %s", chunk) + log.debug( + "Received empty chunk['message'] in streaming chunk: %s", chunk + ) continue token = chunk["message"]["content"] if on_token_callback: diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py index e1088c890..52d624941 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py @@ -52,7 +52,9 @@ def __init__( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), + retry=retry_if_exception_type( + (RateLimitError, APIConnectionError, APITimeoutError) + ), ) def generate( self, @@ -87,7 +89,9 @@ def generate( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), + retry=retry_if_exception_type( + (RateLimitError, APIConnectionError, APITimeoutError) + ), ) async def agenerate( self, @@ -122,7 +126,9 @@ async def agenerate( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), + retry=retry_if_exception_type( + (RateLimitError, APIConnectionError, APITimeoutError) + ), ) def generate_streaming( self, diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py index 3bf481ce2..9886aa0ca 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py @@ -57,7 +57,9 @@ def get_rerank_lists( "top_n": top_n, "documents": documents, } - response = requests.post(url, headers=headers, json=payload, timeout=(1.0, 10.0)) + response = requests.post( + url, headers=headers, json=payload, timeout=(1.0, 10.0) + ) response.raise_for_status() # Raise an error for bad status codes results = response.json()["results"] sorted_docs = [documents[item["index"]] for item in results] diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py index e4a9b550a..da8a9f7b7 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py @@ -58,7 +58,9 @@ def get_rerank_lists( "content-type": Constants.HEADER_CONTENT_TYPE, "authorization": f"Bearer {self.api_key}", } - response = requests.post(url, json=payload, headers=headers, timeout=(1.0, 10.0)) + response = requests.post( + url, json=payload, headers=headers, timeout=(1.0, 10.0) + ) response.raise_for_status() # Raise an error for bad status codes results = response.json()["results"] sorted_docs = [documents[item["index"]] for item in results] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index 47b0f060f..bd2479817 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -59,8 +59,12 @@ def _validate_schema(self, schema: Dict[str, Any]) -> None: check_type(schema, dict, "Input data is not a dictionary.") if "vertexlabels" not in schema or "edgelabels" not in schema: log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") - check_type(schema["vertexlabels"], list, "'vertexlabels' in input data is not a list.") - check_type(schema["edgelabels"], list, "'edgelabels' in input data is not a list.") + check_type( + schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." + ) + check_type( + schema["edgelabels"], list, "'edgelabels' in input data is not a list." + ) def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): property_labels = schema.get("propertykeys", []) @@ -78,13 +82,19 @@ def _process_vertex_labels( for vertex_label in schema["vertexlabels"]: self._validate_vertex_label(vertex_label) properties = vertex_label["properties"] - primary_keys = self._process_keys(vertex_label, "primary_keys", properties[:1]) + primary_keys = self._process_keys( + vertex_label, "primary_keys", properties[:1] + ) if len(primary_keys) == 0: log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") vertex_label["primary_keys"] = primary_keys - nullable_keys = self._process_keys(vertex_label, "nullable_keys", properties[1:]) + nullable_keys = self._process_keys( + vertex_label, "nullable_keys", properties[1:] + ) vertex_label["nullable_keys"] = nullable_keys - self._add_missing_properties(properties, property_labels, property_label_set) + self._add_missing_properties( + properties, property_labels, property_label_set + ) def _process_edge_labels( self, schema: Dict[str, Any], property_labels: list, property_label_set: set @@ -92,13 +102,17 @@ def _process_edge_labels( for edge_label in schema["edgelabels"]: self._validate_edge_label(edge_label) properties = edge_label.get("properties", []) - self._add_missing_properties(properties, property_labels, property_label_set) + self._add_missing_properties( + properties, property_labels, property_label_set + ) def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") if "name" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'name'.") - check_type(vertex_label["name"], str, "'name' in vertex_label is not of correct type.") + check_type( + vertex_label["name"], str, "'name' in vertex_label is not of correct type." + ) if "properties" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'properties'.") check_type( @@ -119,7 +133,9 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: log_and_raise( "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." ) - check_type(edge_label["name"], str, "'name' in edge_label is not of correct type.") + check_type( + edge_label["name"], str, "'name' in edge_label is not of correct type." + ) check_type( edge_label["source_label"], str, @@ -131,9 +147,13 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: "'target_label' in edge_label is not of correct type.", ) - def _process_keys(self, label: Dict[str, Any], key_type: str, default_keys: list) -> list: + def _process_keys( + self, label: Dict[str, Any], key_type: str, default_keys: list + ) -> list: keys = label.get(key_type, default_keys) - check_type(keys, list, f"'{key_type}' in {label['name']} is not of correct type.") + check_type( + keys, list, f"'{key_type}' in {label['name']} is not of correct type." + ) new_keys = [key for key in keys if key in label["properties"]] return new_keys diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py index dc5b15e00..a257ccc4c 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py @@ -138,7 +138,8 @@ def _rerank_with_vertex_degree( if self.method == "bleu": vertex_rerank_res = [ - _bleu_rerank(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list + _bleu_rerank(query, vertex_degree) + [""] + for vertex_degree in vertex_degree_list ] depth = len(vertex_degree_list) @@ -146,11 +147,14 @@ def _rerank_with_vertex_degree( if result not in knowledge_with_degree: knowledge_with_degree[result] = [result] + [""] * (depth - 1) if len(knowledge_with_degree[result]) < depth: - knowledge_with_degree[result] += [""] * (depth - len(knowledge_with_degree[result])) + knowledge_with_degree[result] += [""] * ( + depth - len(knowledge_with_degree[result]) + ) def sort_key(res: str) -> Tuple[int, ...]: return tuple( - vertex_rerank_res[i].index(knowledge_with_degree[res][i]) for i in range(depth) + vertex_rerank_res[i].index(knowledge_with_degree[res][i]) + for i in range(depth) ) sorted_results = sorted(results, key=sort_key) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py index 797ea70ae..c23bf7735 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py @@ -71,7 +71,9 @@ def get_cache_dir() -> str: # Windows (hopefully) else: - local = os.environ.get("LOCALAPPDATA", None) or os.path.expanduser("~\\AppData\\Local") + local = os.environ.get("LOCALAPPDATA", None) or os.path.expanduser( + "~\\AppData\\Local" + ) path = Path(local, "hugegraph_llm") if not os.path.exists(path): diff --git a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py index 0d9967020..0d160e1e5 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py @@ -78,7 +78,11 @@ def _filter_keywords( sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: results.update( - {w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language)} + { + w + for w in sub_tokens + if w not in NLTKHelper().stopwords(lang=self._language) + } ) return list(results) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index 6eb805271..fa6b79f91 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -42,7 +42,9 @@ class RAGPipeline: querying graph databases and vector indices, merging and re-ranking results, and generating answers. """ - def __init__(self, llm: Optional[BaseLLM] = None, embedding: Optional[BaseEmbedding] = None): + def __init__( + self, llm: Optional[BaseLLM] = None, embedding: Optional[BaseEmbedding] = None + ): """ Initialize the RAGPipeline with optional LLM and embedding models. diff --git a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py index b205bb798..52f50fdd6 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py @@ -41,11 +41,15 @@ def clear(self): def example_index_build(self, examples, vector_index: type[VectorStoreBase]): self.operators.append( - BuildGremlinExampleIndex(self.embedding, examples, vector_index=vector_index) + BuildGremlinExampleIndex( + self.embedding, examples, vector_index=vector_index + ) ) return self - def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): + def import_schema( + self, from_hugegraph=None, from_extraction=None, from_user_defined=None + ): if from_hugegraph: self.operators.append(SchemaManager(from_hugegraph)) elif from_user_defined: @@ -57,13 +61,17 @@ def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_def return self def example_index_query(self, num_examples, vector_index: type[VectorStoreBase]): - self.operators.append(GremlinExampleIndexQuery(vector_index, self.embedding, num_examples)) + self.operators.append( + GremlinExampleIndexQuery(vector_index, self.embedding, num_examples) + ) return self def gremlin_generate_synthesize( self, schema, gremlin_prompt: Optional[str] = None, vertices: Optional[List[str]] = None ): - self.operators.append(GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt)) + self.operators.append( + GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt) + ) return self def print_result(self): 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 52626b72b..aa886c0af 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 @@ -50,7 +50,9 @@ def run(self, data: dict) -> Dict[str, Any]: if not schema: # TODO: ensure the function works correctly (update the logic later) self.schema_free_mode(data.get("triples", [])) - log.warning("Using schema_free mode, could try schema_define mode for better effect!") + log.warning( + "Using schema_free mode, could try schema_define mode for better effect!" + ) else: self.init_schema_if_need(schema) self.load_into_graph(vertices, edges, schema) @@ -66,7 +68,9 @@ def _set_default_property(self, key, input_properties, property_label_map): # list or set default_value = [] input_properties[key] = default_value - log.warning("Property '%s' missing in vertex, set to '%s' for now", key, default_value) + log.warning( + "Property '%s' missing in vertex, set to '%s' for now", key, default_value + ) def _handle_graph_creation(self, func, *args, **kwargs): try: @@ -78,11 +82,17 @@ def _handle_graph_creation(self, func, *args, **kwargs): log.error("Error on creating: %s, %s", args, e) return None - def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many-statements + def load_into_graph( + self, vertices, edges, schema + ): # pylint: disable=too-many-statements # pylint: disable=R0912 (too-many-branches) - vertex_label_map = {v_label["name"]: v_label for v_label in schema["vertexlabels"]} + 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"]} + property_label_map = { + p_label["name"]: p_label for p_label in schema["propertykeys"] + } for vertex in vertices: input_label = vertex["label"] @@ -98,7 +108,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- vertex_label = vertex_label_map[input_label] primary_keys = vertex_label["primary_keys"] nullable_keys = vertex_label.get("nullable_keys", []) - non_null_keys = [key for key in vertex_label["properties"] if key not in nullable_keys] + non_null_keys = [ + key for key in vertex_label["properties"] if key not in nullable_keys + ] has_problem = False # 2. Handle primary-keys mode vertex @@ -130,7 +142,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- # 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) + 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(): @@ -167,7 +181,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- 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) + self._handle_graph_creation( + self.client.graph().addEdge, label, start, end, properties + ) def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -191,18 +207,20 @@ def init_schema_if_need(self, schema: dict): source_vertex_label = edge["source_label"] target_vertex_label = edge["target_label"] properties = edge["properties"] - self.schema.edgeLabel(edge_label).sourceLabel(source_vertex_label).targetLabel( - target_vertex_label - ).properties(*properties).nullableKeys(*properties).ifNotExist().create() + self.schema.edgeLabel(edge_label).sourceLabel( + source_vertex_label + ).targetLabel(target_vertex_label).properties(*properties).nullableKeys( + *properties + ).ifNotExist().create() def schema_free_mode(self, 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( - "name" - ).ifNotExist().create() + self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel( + "vertex" + ).properties("name").ifNotExist().create() self.schema.indexLabel("vertexByName").onV("vertex").by( "name" @@ -262,7 +280,9 @@ def _set_property_data_type(self, property_key, data_type): log.warning("UUID type is not supported, use text instead") property_key.asText() else: - log.error("Unknown data type %s for property_key %s", data_type, property_key) + log.error( + "Unknown data type %s for property_key %s", data_type, property_key + ) def _set_property_cardinality(self, property_key, cardinality): if cardinality == PropertyCardinality.SINGLE: @@ -272,9 +292,13 @@ def _set_property_cardinality(self, property_key, cardinality): elif cardinality == PropertyCardinality.SET: property_key.valueSet() else: - log.error("Unknown cardinality %s for property_key %s", cardinality, property_key) + 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: + def _check_property_data_type( + self, data_type: str, cardinality: str, value + ) -> bool: if cardinality in ( PropertyCardinality.LIST.value, PropertyCardinality.SET.value, @@ -304,7 +328,9 @@ def _check_single_data_type(self, data_type: str, value) -> bool: if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below - if data_type == PropertyDataType.DATE.value: # the format should be "yyyy-MM-dd" + if ( + data_type == PropertyDataType.DATE.value + ): # the format should be "yyyy-MM-dd" import re return isinstance(value, str) and re.match(r"^\d{4}-\d{2}-\d{2}$", value) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py index 4c4c167c4..73c9530df 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py @@ -47,5 +47,7 @@ def res = [:]; result = self.graph.gremlin().exec(groovy_code)["data"] if isinstance(result, list) and len(result) > 0: - graph_summary.update({key: result[i].get(key) for i, key in enumerate(keys)}) + graph_summary.update( + {key: result[i].get(key) for i, key in enumerate(keys)} + ) return graph_summary diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py index d9542cd39..52399d99e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py @@ -132,7 +132,9 @@ def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: query_embedding = context.get("query_embedding") self._gremlin_generator.clear() - self._gremlin_generator.example_index_query(num_examples=self._num_gremlin_generate_example) + self._gremlin_generator.example_index_query( + num_examples=self._num_gremlin_generate_example + ) gremlin_response = self._gremlin_generator.gremlin_generate_synthesize( context["simple_schema"], vertices=vertices, gremlin_prompt=self._gremlin_prompt ).run(query=query, query_embedding=query_embedding) @@ -146,7 +148,9 @@ def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: result = self._client.gremlin().exec(gremlin=gremlin)["data"] if result == [None]: result = [] - context["graph_result"] = [json.dumps(item, ensure_ascii=False) for item in result] + context["graph_result"] = [ + json.dumps(item, ensure_ascii=False) for item in result + ] if context["graph_result"]: context["graph_result_flag"] = 1 context["graph_context_head"] = ( @@ -224,7 +228,9 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: "Unable to find vid, downgraded to property query, please confirm if it meets expectation." ) - paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)["data"] + paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)[ + "data" + ] graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( self._format_graph_query_result(query_paths=paths) ) @@ -346,14 +352,18 @@ def _process_vertex( use_id_to_match: bool, v_cache: Set[str], ) -> Tuple[str, int, int]: - matched_str = item["id"] if use_id_to_match else item["props"][self._prop_to_match] + matched_str = ( + item["id"] if use_id_to_match else item["props"][self._prop_to_match] + ) if matched_str in node_cache: flat_rel = flat_rel[:-prior_edge_str_len] return flat_rel, prior_edge_str_len, depth node_cache.add(matched_str) props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v + f"{k}: {self._limit_property_query(v, 'v')}" + for k, v in item["props"].items() + if v ) # TODO: we may remove label id or replace with label name @@ -378,7 +388,9 @@ def _process_edge( e_cache: Set[Tuple[str, str, str]], ) -> Tuple[str, int]: props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v + f"{k}: {self._limit_property_query(v, 'e')}" + for k, v in item["props"].items() + if v ) props_str = f"{{{props_str}}}" if props_str else "" prev_matched_str = ( @@ -395,7 +407,9 @@ def _process_edge( edge_label = item["label"] edge_str = ( - f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" + f"--[{edge_label}]-->" + if item["outV"] == prev_matched_str + else f"<--[{edge_label}]--" ) path_str += edge_str prior_edge_str_len = len(edge_str) @@ -413,14 +427,20 @@ def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: schema = self._get_graph_schema() vertex_props_str, edge_props_str = schema.split("\n")[:2] # TODO: rename to vertex (also need update in the schema) - vertex_props_str = vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") - edge_props_str = edge_props_str[len("Edge properties: ") :].strip("[").strip("]") + vertex_props_str = ( + vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") + ) + edge_props_str = ( + edge_props_str[len("Edge properties: ") :].strip("[").strip("]") + ) vertex_labels = self._extract_label_names(vertex_props_str) edge_labels = self._extract_label_names(edge_props_str) return vertex_labels, edge_labels @staticmethod - def _extract_label_names(source: str, head: str = "name: ", tail: str = ", ") -> List[str]: + def _extract_label_names( + source: str, head: str = "name: ", tail: str = ", " + ) -> List[str]: result = [] for s in source.split(head): end = s.find(tail) @@ -446,7 +466,9 @@ def _get_graph_schema(self, refresh: bool = False) -> str: log.debug("Link(Relation): %s", relationships) return self._schema - def _limit_property_query(self, value: Optional[str], item_type: str) -> Optional[str]: + def _limit_property_query( + self, value: Optional[str], item_type: str + ) -> Optional[str]: # NOTE: we skip the filter for list/set type (e.g., list of string, add it if needed) if not self._limit_property or not isinstance(value, str): return value 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 2f0643a77..e9bccc2f4 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 @@ -40,7 +40,9 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: new_vertex = { - key: vertex[key] for key in ["id", "name", "properties"] if key in vertex + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex } mini_schema["vertexlabels"].append(new_vertex) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index 4f288a9de..d41f99a37 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -62,10 +62,14 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: # !: We have assumed that self.example is not empty queries = [example["query"] for example in self.examples] # TODO: refactor function chain async to avoid blocking - examples_embedding = asyncio.run(get_embeddings_parallel(self.embedding, queries)) + examples_embedding = asyncio.run( + get_embeddings_parallel(self.embedding, queries) + ) embed_dim = len(examples_embedding[0]) if len(self.examples) > 0: - vector_index = self.vector_index.from_name(embed_dim, self.vector_index_name) + vector_index = self.vector_index.from_name( + embed_dim, self.vector_index_name + ) vector_index.add(examples_embedding, self.examples) <<<<<<< HEAD vector_index.to_index_file(self.index_dir, self.filename_prefix) 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 4c46b8906..3bd838168 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 @@ -49,7 +49,9 @@ async def get_embeddings_with_semaphore(vid_list: list[str]) -> Any: None, self.embedding.get_texts_embeddings, vid_list ) - vid_batches = [vids[i : i + batch_size] for i in range(0, len(vids), batch_size)] + vid_batches = [ + vids[i : i + batch_size] for i in range(0, len(vids), batch_size) + ] tasks = [get_embeddings_with_semaphore(batch) for batch in vid_batches] embeddings = [] @@ -62,18 +64,26 @@ async def get_embeddings_with_semaphore(vid_list: list[str]) -> Any: def run(self, context: Dict[str, Any]) -> Dict[str, Any]: vertexlabels = self.sm.schema.getSchema()["vertexlabels"] - all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) + all_pk_flag = all( + data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels + ) past_vids = self.vid_index.get_all_properties() # only support Faiss # TODO: We should build vid vector index separately, especially when the vertices may be very large - present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py + present_vids = context[ + "vertices" + ] # Warning: data truncated by fetch_graph_data.py removed_vids = set(past_vids) - set(present_vids) removed_num = self.vid_index.remove(removed_vids) added_vids = list(set(present_vids) - set(past_vids)) if added_vids: - vids_to_process = self._extract_names(added_vids) if all_pk_flag else added_vids - added_embeddings = asyncio.run(self._get_embeddings_parallel(vids_to_process)) + vids_to_process = ( + self._extract_names(added_vids) if all_pk_flag else added_vids + ) + 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") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index b055cf62e..8d1f3394f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -59,19 +59,23 @@ def __init__( self.embedding.get_embedding_dim(), "gremlin_examples" ) - def _get_match_result(self, context: Dict[str, Any], query: str) -> List[Dict[str, Any]]: + def _get_match_result( + self, context: Dict[str, Any], query: str + ) -> List[Dict[str, Any]]: if self.num_examples <= 0: return [] query_embedding = context.get("query_embedding") if not isinstance(query_embedding, list): query_embedding = self.embedding.get_texts_embeddings([query])[0] - return self.vector_index.search(query_embedding, self.num_examples, dis_threshold=1.8) + return self.vector_index.search( + query_embedding, self.num_examples, dis_threshold=1.8 + ) def _build_default_example_index(self): - properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict( - orient="records" - ) + properties = pd.read_csv( + os.path.join(resource_path, "demo", "text2gremlin.csv") + ).to_dict(orient="records") from concurrent.futures import ThreadPoolExecutor # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py index 49e712d01..923068462 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py @@ -39,7 +39,9 @@ def __init__( topk_per_keyword: int = huge_settings.topk_per_keyword, vector_dis_threshold: float = huge_settings.vector_dis_threshold, ): - self.index_dir = str(os.path.join(resource_path, huge_settings.graph_name, "graph_vids")) + self.index_dir = str( + os.path.join(resource_path, huge_settings.graph_name, "graph_vids") + ) self.vector_index = vector_index.from_name( embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids" ) @@ -65,7 +67,9 @@ def _exact_match_vids(self, keywords: List[str]) -> Tuple[List[str], List[str]]: possible_vids.update([f"{i + 1}:{keyword}" for keyword in keywords]) vids_str = ",".join([f"'{vid}'" for vid in possible_vids]) - resp = self._client.gremlin().exec(SemanticIdQuery.ID_QUERY_TEMPL.format(vids_str=vids_str)) + resp = self._client.gremlin().exec( + SemanticIdQuery.ID_QUERY_TEMPL.format(vids_str=vids_str) + ) searched_vids = [v["id"] for v in resp["data"]] unsearched_keywords = set(keywords) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py index 0b396b77d..cb7d1f1ae 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py @@ -51,7 +51,9 @@ def __init__( self.graph = graph self.result = None - def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): + def import_schema( + self, from_hugegraph=None, from_extraction=None, from_user_defined=None + ): if from_hugegraph: self.operators.append(SchemaManager(from_hugegraph)) elif from_user_defined: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py index 2140abca2..2447c2e06 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py @@ -62,8 +62,10 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( - "\n" + context_str = ( + f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( + "\n" + ) ) final_prompt = self._prompt_template.format( @@ -90,8 +92,12 @@ def init_llm(self, context): if self._question is None: self._question = context.get("query") or None assert self._question is not None, "No question for synthesizing." - context_head_str = context.get("synthesize_context_head") or self._context_head or "" - context_tail_str = context.get("synthesize_context_tail") or self._context_tail or "" + context_head_str = ( + context.get("synthesize_context_head") or self._context_head or "" + ) + context_tail_str = ( + context.get("synthesize_context_tail") or self._context_tail or "" + ) return context_head_str, context_tail_str def handle_vector_graph(self, context): @@ -115,12 +121,16 @@ def handle_vector_graph(self, context): log.warning(graph_result_context) return graph_result_context, vector_result_context - async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[str, Any], None]: + async def run_streaming( + self, context: Dict[str, Any] + ) -> AsyncGenerator[Dict[str, Any], None]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( - "\n" + context_str = ( + f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( + "\n" + ) ) final_prompt = self._prompt_template.format( @@ -153,7 +163,9 @@ async def async_generate( async_tasks = {} if self._raw_answer: final_prompt = self._question - async_tasks["raw_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) + async_tasks["raw_task"] = asyncio.create_task( + self._llm.agenerate(prompt=final_prompt) + ) if self._vector_only_answer: context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip( "\n" @@ -166,8 +178,10 @@ async def async_generate( self._llm.agenerate(prompt=final_prompt) ) if self._graph_only_answer: - context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( - "\n" + context_str = ( + f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( + "\n" + ) ) final_prompt = self._prompt_template.format( @@ -180,7 +194,11 @@ async def async_generate( context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") + context_str = ( + f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip( + "\n" + ) + ) final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -249,8 +267,10 @@ async def async_streaming_generate( ) auto_id += 1 if self._graph_only_answer: - context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( - "\n" + context_str = ( + f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( + "\n" + ) ) final_prompt = self._prompt_template.format( @@ -266,7 +286,11 @@ async def async_streaming_generate( context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") + context_str = ( + f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip( + "\n" + ) + ) final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -292,7 +316,9 @@ async def async_streaming_generate( async_tasks = [asyncio.create_task(anext(gen)) for gen in async_generators] while True: - done, _ = await asyncio.wait(async_tasks, return_when=asyncio.FIRST_COMPLETED) + done, _ = await asyncio.wait( + async_tasks, return_when=asyncio.FIRST_COMPLETED + ) stop_task_num = 0 for task in done: try: @@ -306,7 +332,9 @@ async def async_streaming_generate( break yield context - async def __llm_generate_with_meta_info(self, task_id: int, target_key: str, prompt: str): + async def __llm_generate_with_meta_info( + self, task_id: int, target_key: str, prompt: str + ): # FIXME: Expected type 'AsyncIterable', got 'Coroutine[Any, Any, AsyncGenerator[str, None]]' instead async for token in self._llm.agenerate_streaming(prompt=prompt): yield task_id, target_key, token diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index fd4583263..2c0244d57 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -48,7 +48,9 @@ def _extract_response(self, response: str, label: str = "gremlin") -> str: return match.group(1).strip() return response.strip() - def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional[str]: + def _format_examples( + self, examples: Optional[List[Dict[str, str]]] + ) -> Optional[str]: if not examples: return None example_strings = [] @@ -86,7 +88,9 @@ def _format_properties(self, properties: Optional[List[tuple]]) -> Optional[str] async def async_generate(self, context: Dict[str, Any]): async_tasks = {} query = context.get("query") - raw_example = [{"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"}] + raw_example = [ + {"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"} + ] raw_prompt = self.gremlin_prompt.format( query=query, schema=self.schema, @@ -94,7 +98,9 @@ async def async_generate(self, context: Dict[str, Any]): vertices=self._format_vertices(vertices=self.vertices), properties=self._format_properties(properties=None), ) - async_tasks["raw_answer"] = asyncio.create_task(self.llm.agenerate(prompt=raw_prompt)) + async_tasks["raw_answer"] = asyncio.create_task( + self.llm.agenerate(prompt=raw_prompt) + ) examples = context.get("match_result") init_prompt = self.gremlin_prompt.format( @@ -124,7 +130,9 @@ async def async_generate(self, context: Dict[str, Any]): def sync_generate(self, context: Dict[str, Any]): query = context.get("query") - raw_example = [{"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"}] + raw_example = [ + {"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"} + ] raw_prompt = self.gremlin_prompt.format( query=query, schema=self.schema, 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 a9ac4c050..609c5f22f 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 @@ -197,8 +197,12 @@ def valid(self, element_id: str, max_length: int = 256) -> bool: return True def _filter_long_id(self, graph) -> Dict[str, List[Any]]: - graph["vertices"] = [vertex for vertex in graph["vertices"] if self.valid(vertex["id"])] + graph["vertices"] = [ + vertex for vertex in graph["vertices"] if self.valid(vertex["id"]) + ] graph["edges"] = [ - edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"]) + edge + for edge in graph["edges"] + if self.valid(edge["start"]) and self.valid(edge["end"]) ] return graph diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index 0fa1c0f04..420fc9776 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -113,6 +113,8 @@ def _extract_keywords_from_response( sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: results.update( - w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language) + w + for w in sub_tokens + if w not in NLTKHelper().stopwords(lang=self._language) ) return results diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py index 058d1bce9..a45812393 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py @@ -31,7 +31,9 @@ def __init__(self, llm: BaseLLM): def _load_few_shot_example(self, example_name: str) -> Dict[str, Any]: """Loads and finds the specified few-shot example from the unified JSON file.""" - examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "prompt_examples.json" + ) if not os.path.exists(examples_path): raise FileNotFoundError(f"Examples file not found: {examples_path}") with open(examples_path, "r", encoding="utf-8") as f: @@ -39,7 +41,9 @@ def _load_few_shot_example(self, example_name: str) -> Dict[str, Any]: for example in all_examples: if example.get("name") == example_name: return example - raise ValueError(f"Example with name '{example_name}' not found in prompt_examples.json") + raise ValueError( + f"Example with name '{example_name}' not found in prompt_examples.json" + ) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: """Executes the core logic of prompt generation.""" @@ -48,7 +52,9 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: example_name = context.get("example_name") if not all([source_text, scenario, example_name]): - raise ValueError("Missing required context: source_text, scenario, or example_name.") + raise ValueError( + "Missing required context: source_text, scenario, or example_name." + ) few_shot_example = self._load_few_shot_example(example_name) meta_prompt = prompt_tpl.generate_extract_prompt_template.format( 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 d517db8b9..3ba178c88 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 @@ -67,9 +67,9 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: item_type = item["type"] if item_type == "vertex": label = item["label"] - non_nullable_keys = set(properties_map[item_type][label]["properties"]).difference( - set(properties_map[item_type][label]["nullable_keys"]) - ) + non_nullable_keys = set( + properties_map[item_type][label]["properties"] + ).difference(set(properties_map[item_type][label]["nullable_keys"])) for key in non_nullable_keys: if key not in item["properties"]: item["properties"][key] = "NULL" @@ -82,7 +82,9 @@ 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 + ) -> None: self.llm = llm self.example_prompt = example_prompt self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name @@ -148,7 +150,9 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: and "vertices" in property_graph and "edges" in property_graph ): - log.critical("Invalid property graph format; expecting 'vertices' and 'edges'.") + log.critical( + "Invalid property graph format; expecting 'vertices' and 'edges'." + ) return items # Create sets for valid vertex and edge labels based on the schema @@ -158,7 +162,9 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: def process_items(item_list, valid_labels, item_type): for item in item_list: if not isinstance(item, dict): - log.warning("Invalid property graph item type '%s'.", type(item)) + log.warning( + "Invalid property graph item type '%s'.", type(item) + ) continue if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): log.warning("Invalid item keys '%s'.", item.keys()) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py index 1e33514ca..ae445c206 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py @@ -117,9 +117,13 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: raise ValueError("Context must be a dictionary") if "raw_texts" not in context or not isinstance(context["raw_texts"], list): raise ValueError("'raw_texts' must be a list[str]") - if "query_examples" not in context or not isinstance(context["query_examples"], list): + if "query_examples" not in context or not isinstance( + context["query_examples"], list + ): raise ValueError("'query_examples' must be a list[str]") - if "few_shot_schema" not in context or not isinstance(context["few_shot_schema"], dict): + if "few_shot_schema" not in context or not isinstance( + context["few_shot_schema"], dict + ): raise ValueError("'few_shot_schema' must be a dict") raw_texts = context["raw_texts"] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py index 6beeb0291..98cc97ccf 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py @@ -104,7 +104,9 @@ def nodes_schemas_text_to_list_of_dict(nodes_schemas): properties = json.loads(properties) except json.decoder.JSONDecodeError: properties = {} - result.append({"label": label, "primary_key": primary_key, "properties": properties}) + result.append( + {"label": label, "primary_key": primary_key, "properties": properties} + ) return result @@ -116,7 +118,9 @@ def relationships_schemas_text_to_list_of_dict(relationships_schemas): continue start = relationships_schema_list[0].strip().replace('"', "") end = relationships_schema_list[2].strip().replace('"', "") - relationships_schema_type = relationships_schema_list[1].strip().replace('"', "") + relationships_schema_type = ( + relationships_schema_list[1].strip().replace('"', "") + ) properties = re.search(JSON_REGEX, relationships_schema) if properties is None: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py index 2914c4b28..b5232d268 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py @@ -23,7 +23,9 @@ from hugegraph_llm.utils.log import log -def log_elapsed_time(start_time: float, func: Callable, args: tuple, msg: Optional[str]): +def log_elapsed_time( + start_time: float, func: Callable, args: tuple, msg: Optional[str] +): elapse_time = time.perf_counter() - start_time unit = "s" if elapse_time < 1: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py index b2f485cea..ace3d4b6a 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py @@ -32,7 +32,9 @@ async def _get_batch_with_progress( return result -async def get_embeddings_parallel(embedding: BaseEmbedding, vids: list[str]) -> list[Any]: +async def get_embeddings_parallel( + embedding: BaseEmbedding, vids: list[str] +) -> list[Any]: """Get embeddings for texts in parallel. This function processes text embeddings asynchronously in parallel, using batching and semaphore @@ -60,7 +62,9 @@ async def get_embeddings_parallel(embedding: BaseEmbedding, vids: list[str]) -> embeddings = [] with tqdm(total=len(vid_batches)) as pbar: # Create tasks for each batch with progress bar updates - tasks = [_get_batch_with_progress(embedding, batch, pbar) for batch in vid_batches] + tasks = [ + _get_batch_with_progress(embedding, batch, pbar) for batch in vid_batches + ] # Use asyncio.gather() to preserve order batch_results = await asyncio.gather(*tasks) @@ -74,7 +78,9 @@ async def get_embeddings_parallel(embedding: BaseEmbedding, vids: list[str]) -> def get_filename_prefix(embedding_type: str = None, model_name: str = None) -> str: """Generate filename based on model name.""" - if not (model_name and model_name.strip() and embedding_type and embedding_type.strip()): + if not ( + model_name and model_name.strip() and embedding_type and embedding_type.strip() + ): return "" # Sanitize model_name to prevent path traversal or invalid filename chars safe_embedding_type = embedding_type.replace("/", "_").replace("\\", "_").strip() diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 2f48af332..e7b213c61 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -33,7 +33,9 @@ def get_graph_index_info(): - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) graph_summary_info = builder.fetch_graph_data().run() vector_index = get_vector_index_class(index_settings.cur_vector_index) vector_index_entity = vector_index.from_name( @@ -88,14 +90,18 @@ def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: def extract_graph(input_file, input_text, schema, example_prompt) -> str: texts = read_documents(input_file, input_text) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) if not schema: return "ERROR: please input with correct schema/format." error_message = parse_schema(schema, builder) if error_message: return error_message - builder.chunk_split(texts, "document", "zh").extract_info(example_prompt, "property_graph") + builder.chunk_split(texts, "document", "zh").extract_info( + example_prompt, "property_graph" + ) try: context = builder.run() @@ -122,7 +128,9 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: def update_vid_embedding(): vector_index = get_vector_index_class(index_settings.cur_vector_index) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) builder.fetch_graph_data().build_vertex_id_semantic_index(vector_index) log.debug("Operators: %s", builder.operators) try: @@ -139,7 +147,9 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: data_json = json.loads(data.strip()) log.debug("Import graph data: %s", data) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) if schema: error_message = parse_schema(schema, builder) if error_message: @@ -185,7 +195,9 @@ def build_schema(input_text, query_example, few_shot): except json.JSONDecodeError as e: raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) try: schema = builder.build_schema().run(context) except Exception as e: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index 6da7a6567..082d6bb8c 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -29,7 +29,9 @@ MAX_BACKUP_DIRS = 7 MAX_VERTICES = 100000 MAX_EDGES = 200000 -BACKUP_DIR = str(os.path.join(resource_path, "backup-graph-data-4020", huge_settings.graph_name)) +BACKUP_DIR = str( + os.path.join(resource_path, "backup-graph-data-4020", huge_settings.graph_name) +) def run_gremlin_query(query, fmt=True): @@ -56,21 +58,33 @@ def init_hg_test_data(): schema.vertexLabel("Person").properties( "name", "birthDate" ).useCustomizeStringId().ifNotExist().create() - schema.vertexLabel("Movie").properties("name").useCustomizeStringId().ifNotExist().create() - schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() + schema.vertexLabel("Movie").properties( + "name" + ).useCustomizeStringId().ifNotExist().create() + schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel( + "Movie" + ).ifNotExist().create() - schema.indexLabel("PersonByName").onV("Person").by("name").secondary().ifNotExist().create() - schema.indexLabel("MovieByName").onV("Movie").by("name").secondary().ifNotExist().create() + schema.indexLabel("PersonByName").onV("Person").by( + "name" + ).secondary().ifNotExist().create() + schema.indexLabel("MovieByName").onV("Movie").by( + "name" + ).secondary().ifNotExist().create() graph = client.graph() - graph.addVertex("Person", {"name": "Al Pacino", "birthDate": "1940-04-25"}, id="Al Pacino") + graph.addVertex( + "Person", {"name": "Al Pacino", "birthDate": "1940-04-25"}, id="Al Pacino" + ) graph.addVertex( "Person", {"name": "Robert De Niro", "birthDate": "1943-08-17"}, id="Robert De Niro", ) graph.addVertex("Movie", {"name": "The Godfather"}, id="The Godfather") - graph.addVertex("Movie", {"name": "The Godfather Part II"}, id="The Godfather Part II") + graph.addVertex( + "Movie", {"name": "The Godfather Part II"}, id="The Godfather Part II" + ) graph.addVertex( "Movie", {"name": "The Godfather Coda The Death of Michael Corleone"}, @@ -79,7 +93,9 @@ def init_hg_test_data(): graph.addEdge("ActedIn", "Al Pacino", "The Godfather", {}) graph.addEdge("ActedIn", "Al Pacino", "The Godfather Part II", {}) - graph.addEdge("ActedIn", "Al Pacino", "The Godfather Coda The Death of Michael Corleone", {}) + graph.addEdge( + "ActedIn", "Al Pacino", "The Godfather Coda The Death of Michael Corleone", {} + ) graph.addEdge("ActedIn", "Robert De Niro", "The Godfather Part II", {}) schema.getSchema() graph.close() @@ -118,7 +134,9 @@ def backup_data(): } vertexlabels = client.schema().getSchema()["vertexlabels"] - all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) + all_pk_flag = all( + data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels + ) for filename, query in files.items(): write_backup_file(client, backup_subdir, filename, query, all_pk_flag) @@ -198,7 +216,9 @@ def manage_backup_retention(): # TODO: In the path demo/rag_demo/configs_block.py, # there is a function test_api_connection that is similar to this function, # but it is not straightforward to reuse -def check_graph_db_connection(url: str, name: str, user: str, pwd: str, graph_space: str) -> bool: +def check_graph_db_connection( + url: str, name: str, user: str, pwd: str, graph_space: str +) -> bool: try: if graph_space and graph_space.strip(): test_url = f"{url}/graphspaces/{graph_space}/graphs/{name}/schema" diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 4f56db636..df48ea1e9 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -51,7 +51,9 @@ def read_documents(input_file, input_text): texts.append(text) elif full_path.endswith(".pdf"): # TODO: support PDF file - raise gr.Error("PDF will be supported later! Try to upload text/docx now") + raise gr.Error( + "PDF will be supported later! Try to upload text/docx now" + ) else: raise gr.Error("Please input txt or docx file.") else: @@ -87,8 +89,14 @@ def build_vector_index(input_file, input_text): if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - context = builder.chunk_split(texts, "paragraph", "zh").build_vector_index(vector_index).run() + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) + context = ( + builder.chunk_split(texts, "paragraph", "zh") + .build_vector_index(vector_index) + .run() + ) return json.dumps(context, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/tests/models/embeddings/test_ollama_embedding.py b/hugegraph-llm/src/tests/models/embeddings/test_ollama_embedding.py index a7a9d044c..3abccc268 100644 --- a/hugegraph-llm/src/tests/models/embeddings/test_ollama_embedding.py +++ b/hugegraph-llm/src/tests/models/embeddings/test_ollama_embedding.py @@ -32,5 +32,7 @@ def test_get_cosine_similarity(self): ollama_embedding = OllamaEmbedding(model_name="quentinz/bge-large-zh-v1.5") embedding1 = ollama_embedding.get_text_embedding("hello world") embedding2 = ollama_embedding.get_text_embedding("bye world") - similarity = OllamaEmbedding.similarity(embedding1, embedding2, SimilarityMode.DEFAULT) + similarity = OllamaEmbedding.similarity( + embedding1, embedding2, SimilarityMode.DEFAULT + ) print(similarity) diff --git a/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py b/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py index 317d02879..ad975a37e 100644 --- a/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py +++ b/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py @@ -26,7 +26,9 @@ def setUp(self): def test_schema_check_with_valid_input(self): data = { - "vertexlabels": [{"name": "person", "properties": ["name", "age", "occupation"]}], + "vertexlabels": [ + {"name": "person", "properties": ["name", "age", "occupation"]} + ], "edgelabels": [ { "name": "knows", diff --git a/hugegraph-llm/src/tests/operators/llm_op/test_disambiguate_data.py b/hugegraph-llm/src/tests/operators/llm_op/test_disambiguate_data.py index 04ee42142..e120a1b00 100644 --- a/hugegraph-llm/src/tests/operators/llm_op/test_disambiguate_data.py +++ b/hugegraph-llm/src/tests/operators/llm_op/test_disambiguate_data.py @@ -57,7 +57,9 @@ def setUp(self): "properties": {"name": "www.bob.com", "url": "www.bob.com"}, }, ], - "edges": [{"start": "Alice", "end": "Bob", "type": "roommate", "properties": {}}], + "edges": [ + {"start": "Alice", "end": "Bob", "type": "roommate", "properties": {}} + ], "schema": { "vertices": [ {"vertex_label": "person", "properties": ["name", "age", "occupation"]}, From 08f6857331195ae4a6b4333233c4d31296df2a03 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Wed, 3 Sep 2025 15:59:30 +0800 Subject: [PATCH 29/71] fix --- .../hugegraph_llm/api/models/rag_requests.py | 24 ++--- .../src/hugegraph_llm/api/rag_api.py | 23 ++--- .../src/hugegraph_llm/config/llm_config.py | 16 +--- .../config/models/base_config.py | 12 +-- .../config/models/base_prompt_config.py | 12 +-- .../demo/rag_demo/admin_block.py | 4 +- .../src/hugegraph_llm/demo/rag_demo/app.py | 4 +- .../demo/rag_demo/configs_block.py | 87 +++++-------------- .../demo/rag_demo/other_block.py | 12 +-- .../hugegraph_llm/demo/rag_demo/rag_block.py | 40 +++------ .../demo/rag_demo/text2gremlin_block.py | 34 ++------ .../demo/rag_demo/vector_graph_block.py | 54 ++++-------- .../vector_index/milvus_vector_store.py | 8 +- .../vector_index/qdrant_vector_store.py | 8 +- .../hugegraph_llm/middleware/middleware.py | 4 +- .../models/embeddings/init_embedding.py | 8 +- .../hugegraph_llm/models/embeddings/openai.py | 4 +- .../src/hugegraph_llm/models/llms/init_llm.py | 6 +- .../src/hugegraph_llm/models/llms/ollama.py | 8 +- .../src/hugegraph_llm/models/llms/openai.py | 12 +-- .../hugegraph_llm/models/rerankers/cohere.py | 4 +- .../models/rerankers/siliconflow.py | 4 +- .../operators/common_op/check_schema.py | 40 +++------ .../operators/common_op/merge_dedup_rerank.py | 10 +-- .../operators/common_op/nltk_helper.py | 4 +- .../operators/document_op/word_extract.py | 6 +- .../hugegraph_llm/operators/graph_rag_task.py | 4 +- .../operators/gremlin_generate_task.py | 16 +--- .../hugegraph_op/commit_to_hugegraph.py | 62 ++++--------- .../hugegraph_op/fetch_graph_data.py | 4 +- .../operators/hugegraph_op/graph_rag_query.py | 44 +++------- .../operators/hugegraph_op/schema_manager.py | 4 +- .../index_op/build_gremlin_example_index.py | 8 +- .../index_op/build_semantic_index.py | 20 ++--- .../index_op/gremlin_example_index_query.py | 14 ++- .../operators/index_op/semantic_id_query.py | 8 +- .../operators/kg_construction_task.py | 4 +- .../operators/llm_op/answer_synthesize.py | 60 ++++--------- .../operators/llm_op/gremlin_generate.py | 16 +--- .../operators/llm_op/info_extract.py | 8 +- .../operators/llm_op/keyword_extract.py | 4 +- .../operators/llm_op/prompt_generate.py | 12 +-- .../llm_op/property_graph_extract.py | 18 ++-- .../operators/llm_op/schema_build.py | 8 +- .../llm_op/unstructured_data_utils.py | 8 +- .../src/hugegraph_llm/utils/decorators.py | 4 +- .../hugegraph_llm/utils/embedding_utils.py | 12 +-- .../hugegraph_llm/utils/graph_index_utils.py | 24 ++--- .../hugegraph_llm/utils/hugegraph_utils.py | 40 +++------ .../hugegraph_llm/utils/vector_index_utils.py | 14 +-- .../embeddings/test_ollama_embedding.py | 4 +- .../operators/common_op/test_check_schema.py | 4 +- .../llm_op/test_disambiguate_data.py | 4 +- 53 files changed, 226 insertions(+), 650 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py index 9ebc94eb1..f46aea02c 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py +++ b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py @@ -34,18 +34,12 @@ class GraphConfigRequest(BaseModel): class RAGRequest(BaseModel): query: str = Query(..., description="Query you want to ask") raw_answer: bool = Query(False, description="Use LLM to generate answer directly") - vector_only: bool = Query( - False, description="Use LLM to generate answer with vector" - ) - graph_only: bool = Query( - True, description="Use LLM to generate answer with graph RAG only" - ) + vector_only: bool = Query(False, description="Use LLM to generate answer with vector") + graph_only: bool = Query(True, description="Use LLM to generate answer with graph RAG only") graph_vector_answer: bool = Query( False, description="Use LLM to generate answer with vector & GraphRAG" ) - graph_ratio: float = Query( - 0.5, description="The ratio of GraphRAG ans & vector ans" - ) + graph_ratio: float = Query(0.5, description="The ratio of GraphRAG ans & vector ans") rerank_method: Literal["bleu", "reranker"] = Query( "bleu", description="Method to rerank the results." ) @@ -59,9 +53,7 @@ class RAGRequest(BaseModel): max_graph_items: int = Query( 30, description="Maximum number of items for GQL queries in graph." ) - topk_return_results: int = Query( - 20, description="Number of sorted results to return finally." - ) + topk_return_results: int = Query(20, description="Number of sorted results to return finally.") vector_dis_threshold: float = Query( 0.9, description="Threshold for vector similarity\ @@ -98,9 +90,7 @@ class GraphRAGRequest(BaseModel): max_graph_items: int = Query( 30, description="Maximum number of items for GQL queries in graph." ) - topk_return_results: int = Query( - 20, description="Number of sorted results to return finally." - ) + topk_return_results: int = Query(20, description="Number of sorted results to return finally.") vector_dis_threshold: float = Query( 0.9, description="Threshold for vector similarity \ @@ -115,9 +105,7 @@ class GraphRAGRequest(BaseModel): client_config: Optional[GraphConfigRequest] = Query( None, description="hugegraph server config." ) - get_vertex_only: bool = Query( - False, description="return only keywords & vertex (early stop)." - ) + get_vertex_only: bool = Query(False, description="return only keywords & vertex (early stop).") gremlin_tmpl_num: int = Query( 1, diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py index 1d5b451b1..5c9295efa 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py @@ -67,8 +67,7 @@ def rag_answer_api(req: RAGRequest): # Keep prompt params in the end custom_related_information=req.custom_priority_info, answer_prompt=req.answer_prompt or prompt.answer_prompt, - keywords_extract_prompt=req.keywords_extract_prompt - or prompt.keywords_extract_prompt, + keywords_extract_prompt=req.keywords_extract_prompt or prompt.keywords_extract_prompt, gremlin_prompt=req.gremlin_prompt or prompt.gremlin_generate_prompt, ) # TODO: we need more info in the response for users to understand the query logic @@ -136,9 +135,7 @@ def graph_rag_recall_api(req: GraphRAGRequest): except TypeError as e: log.error("TypeError in graph_rag_recall_api: %s", e) - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) - ) from e + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e except Exception as e: log.error("Unexpected error occurred: %s", e) raise HTTPException( @@ -149,9 +146,7 @@ def graph_rag_recall_api(req: GraphRAGRequest): @router.post("/config/graph", status_code=status.HTTP_201_CREATED) def graph_config_api(req: GraphConfigRequest): # Accept status code - res = apply_graph_conf( - req.url, req.name, req.user, req.pwd, req.gs, origin_call="http" - ) + res = apply_graph_conf(req.url, req.name, req.user, req.pwd, req.gs, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) # TODO: restructure the implement of llm to three types, like "/config/chat_llm" @@ -164,9 +159,7 @@ def llm_config_api(req: LLMConfigRequest): req.api_key, req.api_base, req.language_model, req.max_tokens, origin_call="http" ) else: - res = apply_llm_conf( - req.host, req.port, req.language_model, None, origin_call="http" - ) + res = apply_llm_conf(req.host, req.port, req.language_model, None, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/embedding", status_code=status.HTTP_201_CREATED) @@ -178,9 +171,7 @@ def embedding_config_api(req: LLMConfigRequest): req.api_key, req.api_base, req.language_model, origin_call="http" ) else: - res = apply_embedding_conf( - req.host, req.port, req.language_model, origin_call="http" - ) + res = apply_embedding_conf(req.host, req.port, req.language_model, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/rerank", status_code=status.HTTP_201_CREATED) @@ -192,9 +183,7 @@ def rerank_config_api(req: RerankerConfigRequest): req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http" ) elif req.reranker_type == "siliconflow": - res = apply_reranker_conf( - req.api_key, req.reranker_model, None, origin_call="http" - ) + res = apply_reranker_conf(req.api_key, req.reranker_model, None, origin_call="http") else: res = status.HTTP_501_NOT_IMPLEMENTED return generate_response(RAGResponse(status_code=res, message="Missing Value")) diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index 8b4f274ee..af22b71b0 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -32,19 +32,13 @@ class LLMConfig(BaseConfig): embedding_type: Optional[Literal["openai", "litellm", "ollama/local"]] = "openai" reranker_type: Optional[Literal["cohere", "siliconflow"]] = None # 1. OpenAI settings - openai_chat_api_base: str = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) + openai_chat_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") openai_chat_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_chat_language_model: str = "gpt-4.1-mini" - openai_extract_api_base: str = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) + openai_extract_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") openai_extract_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_extract_language_model: str = "gpt-4.1-mini" - openai_text2gql_api_base: str = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) + openai_text2gql_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") openai_text2gql_api_key: str | None = os.environ.get("OPENAI_API_KEY") openai_text2gql_language_model: str = "gpt-4.1-mini" openai_embedding_api_base: str = os.environ.get( @@ -57,9 +51,7 @@ class LLMConfig(BaseConfig): openai_extract_tokens: int = 256 openai_text2gql_tokens: int = 4096 # 2. Rerank settings - cohere_base_url: str = os.environ.get( - "CO_API_URL", "https://api.cohere.com/v1/rerank" - ) + cohere_base_url: str = os.environ.get("CO_API_URL", "https://api.cohere.com/v1/rerank") reranker_api_key: Optional[str] = None reranker_model: Optional[str] = None # 3. Ollama settings diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py index 4ec9256c5..5fec3a778 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py @@ -24,9 +24,7 @@ from hugegraph_llm.utils.log import log dir_name = os.path.dirname -env_path = os.path.join( - os.getcwd(), ".env" -) # Load .env from the current working directory +env_path = os.path.join(os.getcwd(), ".env") # Load .env from the current working directory class BaseConfig(BaseSettings): @@ -90,9 +88,7 @@ def check_env(self): # Step 2: Add missing config items to .env self._sync_object_to_env(env_config, config_dict) except Exception as e: - log.error( - "An error occurred when checking the .env variable file: %s", str(e) - ) + log.error("An error occurred when checking the .env variable file: %s", str(e)) raise def _sync_env_to_object(self, env_config, config_dict): @@ -143,9 +139,7 @@ def __init__(self, **data): # Synchronize configurations between the object and .env file self.check_env() - log.info( - "The %s file was loaded. Class: %s", env_path, self.__class__.__name__ - ) + log.info("The %s file was loaded. Class: %s", env_path, self.__class__.__name__) except Exception as e: log.error("An error occurred when initializing the configuration object: %s", str(e)) raise diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py index b15bad0a6..2369d01a6 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py @@ -57,9 +57,7 @@ def ensure_yaml_file_exists(self): current_dir = Path.cwd().resolve() project_root = get_project_root() if current_dir == project_root: - log.info( - "Current working directory is the project root, proceeding to run the app." - ) + log.info("Current working directory is the project root, proceeding to run the app.") else: error_msg = ( f"Current working directory is not the project root. " @@ -124,9 +122,7 @@ def to_literal(val): "gremlin_generate_prompt": to_literal(self.gremlin_generate_prompt), "doc_input_text": to_literal(self.doc_input_text), "_language_generated": str(self.llm_settings.language).lower().strip(), - "generate_extract_prompt_template": to_literal( - self.generate_extract_prompt_template - ), + "generate_extract_prompt_template": to_literal(self.generate_extract_prompt_template), } with open(yaml_file_path, "w", encoding="utf-8") as file: yaml.dump(data, file, allow_unicode=True, sort_keys=False, default_flow_style=False) @@ -154,9 +150,7 @@ def generate_yaml_file(self): self.keywords_extract_prompt = self.keywords_extract_prompt_EN self.doc_input_text = self.doc_input_text_EN self.save_to_yaml() - log.info( - "Prompt file '%s' has been generated with default values.", yaml_file_path - ) + log.info("Prompt file '%s' has been generated with default values.", yaml_file_path) def update_yaml_file(self): self.save_to_yaml() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index d3beebcbb..1a157d387 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -143,9 +143,7 @@ def create_admin_block(): with gr.Row(): with gr.Column(): # Button to clear LLM Server log, initially hidden - clear_llm_server_button = gr.Button( - "Clear LLM Server Log", visible=False - ) + clear_llm_server_button = gr.Button("Clear LLM Server Log", visible=False) with gr.Column(): # Button to refresh LLM Server log manually refresh_llm_server_button = gr.Button( 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 a78e62361..b0979763c 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -166,9 +166,7 @@ def create_app(): prompt.update_yaml_file() assert admin_settings.enable_login auth_enabled = admin_settings.enable_login.lower() == "true" - log.info( - "(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled" - ) + log.info("(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled") api_auth = APIRouter(dependencies=[Depends(authenticate)] if auth_enabled else []) hugegraph_llm = init_rag_ui() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 721cf272c..112892e7c 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -71,9 +71,7 @@ def test_api_connection( log.debug("Request URL: %s", url) try: if method.upper() == "GET": - resp = requests.get( - url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth - ) + resp = requests.get(url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth) elif method.upper() == "POST": resp = requests.post( url, @@ -108,9 +106,7 @@ def test_api_connection( return resp.status_code -def config_qianfan_model( - arg1, arg2, arg3=None, settings_prefix=None, origin_call=None -) -> int: +def config_qianfan_model(arg1, arg2, arg3=None, settings_prefix=None, origin_call=None) -> int: setattr(llm_settings, f"qianfan_{settings_prefix}_api_key", arg1) setattr(llm_settings, f"qianfan_{settings_prefix}_secret_key", arg2) if arg3: @@ -146,9 +142,7 @@ def apply_embedding_config(arg1, arg2, arg3, origin_call=None) -> int: llm_settings.ollama_embedding_host = arg1 llm_settings.ollama_embedding_port = int(arg2) llm_settings.ollama_embedding_model = arg3 - status_code = test_api_connection( - f"http://{arg1}:{arg2}", origin_call=origin_call - ) + status_code = test_api_connection(f"http://{arg1}:{arg2}", origin_call=origin_call) elif embedding_option == "litellm": llm_settings.litellm_embedding_api_key = arg1 llm_settings.litellm_embedding_api_base = arg2 @@ -239,8 +233,7 @@ def apply_llm_config( setattr(llm_settings, f"openai_{current_llm_config}_tokens", int(max_tokens)) test_url = ( - getattr(llm_settings, f"openai_{current_llm_config}_api_base") - + "/chat/completions" + getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" ) data = { "model": model_name, @@ -254,9 +247,7 @@ def apply_llm_config( elif llm_option == "ollama/local": setattr(llm_settings, f"ollama_{current_llm_config}_host", api_key_or_host) - setattr( - llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port) - ) + setattr(llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port)) setattr(llm_settings, f"ollama_{current_llm_config}_language_model", model_name) status_code = test_api_connection( f"http://{api_key_or_host}:{api_base_or_port}", origin_call=origin_call @@ -264,12 +255,8 @@ def apply_llm_config( elif llm_option == "litellm": setattr(llm_settings, f"litellm_{current_llm_config}_api_key", api_key_or_host) - setattr( - llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port - ) - setattr( - llm_settings, f"litellm_{current_llm_config}_language_model", model_name - ) + setattr(llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port) + setattr(llm_settings, f"litellm_{current_llm_config}_language_model", model_name) setattr(llm_settings, f"litellm_{current_llm_config}_tokens", int(max_tokens)) status_code = test_litellm_chat( @@ -396,13 +383,9 @@ def chat_llm_settings(llm_type): ), ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_chat_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_chat_op, inputs=llm_config_input) # Determine whether there are Settings in the.env file env_path = os.path.join( os.getcwd(), ".env" @@ -442,9 +425,7 @@ def extract_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr( - llm_settings, "openai_extract_language_model" - ), + value=getattr(llm_settings, "openai_extract_language_model"), label="model_name", ), gr.Textbox( @@ -463,9 +444,7 @@ def extract_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr( - llm_settings, "ollama_extract_language_model" - ), + value=getattr(llm_settings, "ollama_extract_language_model"), label="model_name", ), gr.Textbox(value="", visible=False), @@ -483,9 +462,7 @@ def extract_llm_settings(llm_type): info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr( - llm_settings, "litellm_extract_language_model" - ), + value=getattr(llm_settings, "litellm_extract_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -495,13 +472,9 @@ def extract_llm_settings(llm_type): ), ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_extract_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) with gr.Tab(label="text2gql"): text2gql_llm_dropdown = gr.Dropdown( @@ -526,9 +499,7 @@ def text2gql_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr( - llm_settings, "openai_text2gql_language_model" - ), + value=getattr(llm_settings, "openai_text2gql_language_model"), label="model_name", ), gr.Textbox( @@ -547,9 +518,7 @@ def text2gql_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr( - llm_settings, "ollama_text2gql_language_model" - ), + value=getattr(llm_settings, "ollama_text2gql_language_model"), label="model_name", ), gr.Textbox(value="", visible=False), @@ -567,9 +536,7 @@ def text2gql_llm_settings(llm_type): info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr( - llm_settings, "litellm_text2gql_language_model" - ), + value=getattr(llm_settings, "litellm_text2gql_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -579,13 +546,9 @@ def text2gql_llm_settings(llm_type): ), ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_text2gql_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_text2gql_op, inputs=llm_config_input) with gr.Accordion("3. Set up the Embedding.", open=False): embedding_dropdown = gr.Dropdown( @@ -672,9 +635,7 @@ def embedding_settings(embedding_type): @gr.render(inputs=[reranker_dropdown]) def reranker_settings(reranker_type): - llm_settings.reranker_type = ( - reranker_type if reranker_type != "None" else None - ) + llm_settings.reranker_type = reranker_type if reranker_type != "None" else None if reranker_type == "cohere": with gr.Row(): reranker_config_input = [ @@ -683,12 +644,8 @@ def reranker_settings(reranker_type): label="api_key", type="password", ), - gr.Textbox( - value=lambda: llm_settings.reranker_model, label="model" - ), - gr.Textbox( - value=lambda: llm_settings.cohere_base_url, label="base_url" - ), + gr.Textbox(value=lambda: llm_settings.reranker_model, label="model"), + gr.Textbox(value=lambda: llm_settings.cohere_base_url, label="base_url"), ] elif reranker_type == "siliconflow": with gr.Row(): diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py index 3f8089b6d..82650b907 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py @@ -34,13 +34,9 @@ def create_other_block(): inp = gr.Textbox( value="g.V().limit(10)", label="Gremlin query", show_copy_button=True, lines=8 ) - out = gr.Code( - label="Output", language="json", elem_classes="code-container-show" - ) + out = gr.Code(label="Output", language="json", elem_classes="code-container-show") btn = gr.Button("Run Gremlin query") - btn.click( - fn=run_gremlin_query, inputs=[inp], outputs=out - ) # pylint: disable=no-member + btn.click(fn=run_gremlin_query, inputs=[inp], outputs=out) # pylint: disable=no-member gr.Markdown("---") with gr.Row(): @@ -55,9 +51,7 @@ def create_other_block(): inp = [] out = gr.Textbox(label="Init Graph Demo Result", show_copy_button=True) btn = gr.Button("(BETA) Init HugeGraph test data (🚧)") - btn.click( - fn=init_hg_test_data, inputs=inp, outputs=out - ) # pylint: disable=no-member + btn.click(fn=init_hg_test_data, inputs=inp, outputs=out) # pylint: disable=no-member @asynccontextmanager diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 31a2d3b09..2a817ada4 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -114,9 +114,7 @@ def rag_answer( max_graph_items=max_graph_items, ) if context.get("switch_to_bleu"): - gr.Warning( - "Online reranker fails, automatically switches to local bleu rerank." - ) + gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") return ( context.get("raw_answer", ""), context.get("vector_only_answer", ""), @@ -224,9 +222,7 @@ async def rag_answer_streaming( graph_search=graph_search, ) if context.get("switch_to_bleu"): - gr.Warning( - "Online reranker fails, automatically switches to local bleu rerank." - ) + gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") answer_synthesize = AnswerSynthesize( raw_answer=raw_answer, vector_only_answer=vector_only_answer, @@ -236,9 +232,7 @@ async def rag_answer_streaming( ) async for context in answer_synthesize.run_streaming(context): if context.get("switch_to_bleu"): - gr.Warning( - "Online reranker fails, automatically switches to local bleu rerank." - ) + gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") yield ( context.get("raw_answer", ""), context.get("vector_only_answer", ""), @@ -308,9 +302,7 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): - raw_radio = gr.Radio( - choices=[True, False], value=False, label="Basic LLM Answer" - ) + raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") vector_only_radio = gr.Radio( choices=[True, False], value=False, label="Vector-only Answer" ) @@ -395,9 +387,7 @@ def toggle_slider(enable): # FIXME: "demo" might conflict with the graph name, it should be modified. answers_path = os.path.join(resource_path, "demo", "questions_answers.xlsx") questions_path = os.path.join(resource_path, "demo", "questions.xlsx") - questions_template_path = os.path.join( - resource_path, "demo", "questions_template.xlsx" - ) + questions_template_path = os.path.join(resource_path, "demo", "questions_template.xlsx") def read_file_to_excel(file: Any, line_count: Optional[int] = None): df = None @@ -477,18 +467,12 @@ def several_rag_answer( file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" ) with gr.Column(): - test_template_file = os.path.join( - resource_path, "demo", "questions_template.xlsx" - ) + test_template_file = os.path.join(resource_path, "demo", "questions_template.xlsx") gr.File(value=test_template_file, label="Download Template File") - answer_max_line_count = gr.Number( - 1, label="Max Lines To Show", minimum=1, maximum=40 - ) + answer_max_line_count = gr.Number(1, label="Max Lines To Show", minimum=1, maximum=40) answers_btn = gr.Button("Generate Answer (Batch)", variant="primary") # TODO: Set individual progress bars for dataframe - qa_dataframe = gr.DataFrame( - label="Questions & Answers (Preview)", headers=tests_df_headers - ) + qa_dataframe = gr.DataFrame(label="Questions & Answers (Preview)", headers=tests_df_headers) answers_btn.click( several_rag_answer, inputs=[ @@ -506,12 +490,8 @@ def several_rag_answer( ], outputs=[qa_dataframe, gr.File(label="Download Answered File", min_width=40)], ) - questions_file.change( - read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count] - ) - answer_max_line_count.change( - change_showing_excel, answer_max_line_count, qa_dataframe - ) + questions_file.change(read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count]) + answer_max_line_count.change(change_showing_excel, answer_max_line_count, qa_dataframe) return ( inp, answer_prompt_input, diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index af55cd86f..8fbb01c25 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -96,12 +96,8 @@ def build_example_vector_index(temp_file) -> dict: timestamp = datetime.now().strftime("%Y%m%d%H%M%S") _, file_name = os.path.split(f"{name}_{timestamp}{ext}") log.info("Copying file to: %s", file_name) - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - target_file = os.path.join( - resource_path, folder_name, "gremlin_examples", file_name - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + target_file = os.path.join(resource_path, folder_name, "gremlin_examples", file_name) try: import shutil @@ -201,13 +197,9 @@ def gremlin_generate( processed_schema, short_schema = _process_schema(schema, generator, sm) if processed_schema is None and short_schema is None: - return GremlinResult.error( - "Invalid JSON schema, please check the format carefully." - ) + return GremlinResult.error("Invalid JSON schema, please check the format carefully.") - updated_schema = ( - sm.simple_schema(processed_schema) if short_schema else processed_schema - ) + updated_schema = sm.simple_schema(processed_schema) if short_schema else processed_schema store_schema(str(updated_schema), inp, gremlin_prompt) output_types = _configure_output_types(requested_outputs) @@ -239,11 +231,7 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = { - key: vertex[key] - for key in ["id", "name", "properties"] - if key in vertex - } + new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) @@ -294,9 +282,7 @@ def create_text2gremlin_block() -> Tuple: with gr.Row(): btn = gr.Button("Build Example Vector Index", variant="primary") - btn.click( - build_example_vector_index, inputs=[file], outputs=[out] - ) # pylint: disable=no-member + btn.click(build_example_vector_index, inputs=[file], outputs=[out]) # pylint: disable=no-member gr.Markdown("## Nature Language To Gremlin") with gr.Row(): @@ -311,12 +297,8 @@ def create_text2gremlin_block() -> Tuple: language="javascript", elem_classes="code-container-show", ) - initialized_out = gr.Textbox( - label="Gremlin With Template", show_copy_button=True - ) - raw_out = gr.Textbox( - label="Gremlin Without Template", show_copy_button=True - ) + initialized_out = gr.Textbox(label="Gremlin With Template", show_copy_button=True) + raw_out = gr.Textbox(label="Gremlin Without Template", show_copy_button=True) tmpl_exec_out = gr.Code( label="Query With Template Output", language="json", diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index c5f47d0b5..5f63e4964 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -61,9 +61,7 @@ def generate_prompt_for_ui(source_text, scenario, example_name): Handles the UI logic for generating a new prompt. It calls the PromptGenerate operator. """ if not all([source_text, scenario, example_name]): - gr.Warning( - "Please provide original text, expected scenario, and select an example!" - ) + gr.Warning("Please provide original text, expected scenario, and select an example!") return gr.update() try: prompt_generator = PromptGenerate(llm=LLMs().get_chat_llm()) @@ -87,9 +85,7 @@ def generate_prompt_for_ui(source_text, scenario, example_name): def load_example_names(): """Load all candidate examples""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "prompt_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return [example.get("name", "Unnamed example") for example in examples] @@ -110,22 +106,16 @@ def load_query_examples(): ), ) if language.upper() == "CN": - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples_CN.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples_CN.json") else: - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) except (FileNotFoundError, json.JSONDecodeError): try: - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -136,9 +126,7 @@ def load_query_examples(): def load_schema_fewshot_examples(): """Load few-shot examples from a JSON file""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "schema_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "schema_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -149,14 +137,10 @@ def load_schema_fewshot_examples(): def update_example_preview(example_name): """Update the display content based on the selected example name.""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "prompt_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") with open(examples_path, "r", encoding="utf-8") as f: all_examples = json.load(f) - selected_example = next( - (ex for ex in all_examples if ex.get("name") == example_name), None - ) + selected_example = next((ex for ex in all_examples if ex.get("name") == example_name), None) if selected_example: return ( @@ -201,9 +185,7 @@ def _create_prompt_helper_block(demo, input_text, info_extract_template): interactive=False, ) - generate_prompt_btn = gr.Button( - "🚀 Auto-generate Graph Extract Prompt", variant="primary" - ) + generate_prompt_btn = gr.Button("🚀 Auto-generate Graph Extract Prompt", variant="primary") # Bind the change event of the dropdown menu few_shot_dropdown.change( fn=update_example_preview, @@ -295,9 +277,7 @@ def create_vector_graph_block(): lines=15, max_lines=29, ) - out = gr.Code( - label="Output Info", language="json", elem_classes="code-container-edit" - ) + out = gr.Code(label="Output Info", language="json", elem_classes="code-container-edit") with gr.Row(): with gr.Accordion("Get RAG Info", open=False): @@ -306,12 +286,8 @@ def create_vector_graph_block(): graph_index_btn0 = gr.Button("Get Graph Index Info", size="sm") with gr.Accordion("Clear RAG Data", open=False): with gr.Column(): - vector_index_btn1 = gr.Button( - "Clear Chunks Vector Index", size="sm" - ) - graph_index_btn1 = gr.Button( - "Clear Graph Vid Vector Index", size="sm" - ) + vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") + graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") vector_import_bt = gr.Button("Import into Vector", variant="primary") @@ -383,9 +359,9 @@ def create_vector_graph_block(): inputs=[input_text, input_schema, info_extract_template], ) - graph_loading_bt.click( - import_graph_data, inputs=[out, input_schema], outputs=[out] - ).then(update_vid_embedding).then( + graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( + update_vid_embedding + ).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py index fae5ca860..b168b8c39 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -73,15 +73,11 @@ def __init__( def _create_collection(self): """Create a new collection in Milvus.""" - id_field = FieldSchema( - name="id", dtype=DataType.INT64, is_primary=True, auto_id=True - ) + id_field = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True) vector_field = FieldSchema( name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embed_dim ) - property_field = FieldSchema( - name="property", dtype=DataType.VARCHAR, max_length=65535 - ) + property_field = FieldSchema(name="property", dtype=DataType.VARCHAR, max_length=65535) original_id_field = FieldSchema(name="original_id", dtype=DataType.INT64) schema = CollectionSchema( diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index 14b97fbff..ca7761ecd 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -28,9 +28,7 @@ class QdrantVectorIndex(VectorStoreBase): - def __init__( - self, name: str, host: str, port: int, api_key=None, embed_dim: int = 1024 - ): + def __init__(self, name: str, host: str, port: int, api_key=None, embed_dim: int = 1024): self.embed_dim = embed_dim self.host = host self.port = port @@ -117,9 +115,7 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: return remove_num - def search( - self, query_vector: List[float], top_k: int = 5, dis_threshold: float = 0.9 - ): + def search(self, query_vector: List[float], top_k: int = 5, dis_threshold: float = 0.9): search_result = self.client.search( collection_name=self.name, query_vector=query_vector, limit=top_k ) diff --git a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py index 5d98ebdf8..f13e11e6f 100644 --- a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py +++ b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py @@ -41,9 +41,7 @@ async def dispatch(self, request: Request, call_next): unit = "s" response.headers["X-Process-Time"] = f"{process_time:.2f} {unit}" - log.info( - "Request process time: %.2f ms, code=%d", process_time, response.status_code - ) + log.info("Request process time: %.2f ms, code=%d", process_time, response.status_code) log.info( "%s - Args: %s, IP: %s, URL: %s", request.method, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index 8e6af2774..d96840911 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -65,9 +65,7 @@ def __init__(self): def get_embedding(self): if self.embedding_type == "openai": - assert ( - llm_settings.openai_embedding_model_dim - ), "openai_embedding_model_dim is need" + assert llm_settings.openai_embedding_model_dim, "openai_embedding_model_dim is need" return OpenAIEmbedding( embedding_dimension=llm_settings.openai_embedding_model_dim, model_name=llm_settings.openai_embedding_model, @@ -75,9 +73,7 @@ def get_embedding(self): api_base=llm_settings.openai_embedding_api_base, ) if self.embedding_type == "ollama/local": - assert ( - llm_settings.ollama_embedding_model_dim - ), "ollama_embedding_model_dim is need" + assert llm_settings.ollama_embedding_model_dim, "ollama_embedding_model_dim is need" return OllamaEmbedding( <<<<<<< HEAD model_name=llm_settings.ollama_embedding_model, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index 15d928286..6e30cca71 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -91,7 +91,5 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] A list of embedding vectors, where each vector is a list of floats. The order of embeddings should match the order of input texts. """ - response = await self.aclient.embeddings.create( - input=texts, model=self.model_name - ) + response = await self.aclient.embeddings.create(input=texts, model=self.model_name) return [data.embedding for data in response.data] 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 7e1eaab68..9121fca09 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py @@ -173,8 +173,4 @@ def get_text2gql_llm(self): if __name__ == "__main__": client = LLMs().get_chat_llm() print(client.generate(prompt="What is the capital of China?")) - print( - client.generate( - messages=[{"role": "user", "content": "What is the capital of China?"}] - ) - ) + print(client.generate(messages=[{"role": "user", "content": "What is the capital of China?"}])) diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py index c15c5440e..199384b12 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py @@ -29,9 +29,7 @@ class OllamaClient(BaseLLM): """LLM wrapper should take in a prompt and return a string.""" - def __init__( - self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwargs - ): + def __init__(self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwargs): self.model = model self.client = ollama.Client(host=f"http://{host}:{port}", **kwargs) self.async_client = ollama.AsyncClient(host=f"http://{host}:{port}", **kwargs) @@ -101,9 +99,7 @@ def generate_streaming( for chunk in self.client.chat(model=self.model, messages=messages, stream=True): if not chunk["message"]: - log.debug( - "Received empty chunk['message'] in streaming chunk: %s", chunk - ) + log.debug("Received empty chunk['message'] in streaming chunk: %s", chunk) continue token = chunk["message"]["content"] if on_token_callback: diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py index 52d624941..e1088c890 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py @@ -52,9 +52,7 @@ def __init__( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type( - (RateLimitError, APIConnectionError, APITimeoutError) - ), + retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), ) def generate( self, @@ -89,9 +87,7 @@ def generate( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type( - (RateLimitError, APIConnectionError, APITimeoutError) - ), + retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), ) async def agenerate( self, @@ -126,9 +122,7 @@ async def agenerate( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception_type( - (RateLimitError, APIConnectionError, APITimeoutError) - ), + retry=retry_if_exception_type((RateLimitError, APIConnectionError, APITimeoutError)), ) def generate_streaming( self, diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py index 9886aa0ca..3bf481ce2 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py @@ -57,9 +57,7 @@ def get_rerank_lists( "top_n": top_n, "documents": documents, } - response = requests.post( - url, headers=headers, json=payload, timeout=(1.0, 10.0) - ) + response = requests.post(url, headers=headers, json=payload, timeout=(1.0, 10.0)) response.raise_for_status() # Raise an error for bad status codes results = response.json()["results"] sorted_docs = [documents[item["index"]] for item in results] diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py index da8a9f7b7..e4a9b550a 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py @@ -58,9 +58,7 @@ def get_rerank_lists( "content-type": Constants.HEADER_CONTENT_TYPE, "authorization": f"Bearer {self.api_key}", } - response = requests.post( - url, json=payload, headers=headers, timeout=(1.0, 10.0) - ) + response = requests.post(url, json=payload, headers=headers, timeout=(1.0, 10.0)) response.raise_for_status() # Raise an error for bad status codes results = response.json()["results"] sorted_docs = [documents[item["index"]] for item in results] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index bd2479817..47b0f060f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -59,12 +59,8 @@ def _validate_schema(self, schema: Dict[str, Any]) -> None: check_type(schema, dict, "Input data is not a dictionary.") if "vertexlabels" not in schema or "edgelabels" not in schema: log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") - check_type( - schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." - ) - check_type( - schema["edgelabels"], list, "'edgelabels' in input data is not a list." - ) + check_type(schema["vertexlabels"], list, "'vertexlabels' in input data is not a list.") + check_type(schema["edgelabels"], list, "'edgelabels' in input data is not a list.") def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): property_labels = schema.get("propertykeys", []) @@ -82,19 +78,13 @@ def _process_vertex_labels( for vertex_label in schema["vertexlabels"]: self._validate_vertex_label(vertex_label) properties = vertex_label["properties"] - primary_keys = self._process_keys( - vertex_label, "primary_keys", properties[:1] - ) + primary_keys = self._process_keys(vertex_label, "primary_keys", properties[:1]) if len(primary_keys) == 0: log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") vertex_label["primary_keys"] = primary_keys - nullable_keys = self._process_keys( - vertex_label, "nullable_keys", properties[1:] - ) + nullable_keys = self._process_keys(vertex_label, "nullable_keys", properties[1:]) vertex_label["nullable_keys"] = nullable_keys - self._add_missing_properties( - properties, property_labels, property_label_set - ) + self._add_missing_properties(properties, property_labels, property_label_set) def _process_edge_labels( self, schema: Dict[str, Any], property_labels: list, property_label_set: set @@ -102,17 +92,13 @@ def _process_edge_labels( for edge_label in schema["edgelabels"]: self._validate_edge_label(edge_label) properties = edge_label.get("properties", []) - self._add_missing_properties( - properties, property_labels, property_label_set - ) + self._add_missing_properties(properties, property_labels, property_label_set) def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") if "name" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'name'.") - check_type( - vertex_label["name"], str, "'name' in vertex_label is not of correct type." - ) + check_type(vertex_label["name"], str, "'name' in vertex_label is not of correct type.") if "properties" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'properties'.") check_type( @@ -133,9 +119,7 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: log_and_raise( "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." ) - check_type( - edge_label["name"], str, "'name' in edge_label is not of correct type." - ) + check_type(edge_label["name"], str, "'name' in edge_label is not of correct type.") check_type( edge_label["source_label"], str, @@ -147,13 +131,9 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: "'target_label' in edge_label is not of correct type.", ) - def _process_keys( - self, label: Dict[str, Any], key_type: str, default_keys: list - ) -> list: + def _process_keys(self, label: Dict[str, Any], key_type: str, default_keys: list) -> list: keys = label.get(key_type, default_keys) - check_type( - keys, list, f"'{key_type}' in {label['name']} is not of correct type." - ) + check_type(keys, list, f"'{key_type}' in {label['name']} is not of correct type.") new_keys = [key for key in keys if key in label["properties"]] return new_keys diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py index a257ccc4c..dc5b15e00 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py @@ -138,8 +138,7 @@ def _rerank_with_vertex_degree( if self.method == "bleu": vertex_rerank_res = [ - _bleu_rerank(query, vertex_degree) + [""] - for vertex_degree in vertex_degree_list + _bleu_rerank(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list ] depth = len(vertex_degree_list) @@ -147,14 +146,11 @@ def _rerank_with_vertex_degree( if result not in knowledge_with_degree: knowledge_with_degree[result] = [result] + [""] * (depth - 1) if len(knowledge_with_degree[result]) < depth: - knowledge_with_degree[result] += [""] * ( - depth - len(knowledge_with_degree[result]) - ) + knowledge_with_degree[result] += [""] * (depth - len(knowledge_with_degree[result])) def sort_key(res: str) -> Tuple[int, ...]: return tuple( - vertex_rerank_res[i].index(knowledge_with_degree[res][i]) - for i in range(depth) + vertex_rerank_res[i].index(knowledge_with_degree[res][i]) for i in range(depth) ) sorted_results = sorted(results, key=sort_key) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py index c23bf7735..797ea70ae 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/nltk_helper.py @@ -71,9 +71,7 @@ def get_cache_dir() -> str: # Windows (hopefully) else: - local = os.environ.get("LOCALAPPDATA", None) or os.path.expanduser( - "~\\AppData\\Local" - ) + local = os.environ.get("LOCALAPPDATA", None) or os.path.expanduser("~\\AppData\\Local") path = Path(local, "hugegraph_llm") if not os.path.exists(path): diff --git a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py index 0d160e1e5..0d9967020 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py @@ -78,11 +78,7 @@ def _filter_keywords( sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: results.update( - { - w - for w in sub_tokens - if w not in NLTKHelper().stopwords(lang=self._language) - } + {w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language)} ) return list(results) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index fa6b79f91..6eb805271 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -42,9 +42,7 @@ class RAGPipeline: querying graph databases and vector indices, merging and re-ranking results, and generating answers. """ - def __init__( - self, llm: Optional[BaseLLM] = None, embedding: Optional[BaseEmbedding] = None - ): + def __init__(self, llm: Optional[BaseLLM] = None, embedding: Optional[BaseEmbedding] = None): """ Initialize the RAGPipeline with optional LLM and embedding models. diff --git a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py index 52f50fdd6..b205bb798 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py @@ -41,15 +41,11 @@ def clear(self): def example_index_build(self, examples, vector_index: type[VectorStoreBase]): self.operators.append( - BuildGremlinExampleIndex( - self.embedding, examples, vector_index=vector_index - ) + BuildGremlinExampleIndex(self.embedding, examples, vector_index=vector_index) ) return self - def import_schema( - self, from_hugegraph=None, from_extraction=None, from_user_defined=None - ): + def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): if from_hugegraph: self.operators.append(SchemaManager(from_hugegraph)) elif from_user_defined: @@ -61,17 +57,13 @@ def import_schema( return self def example_index_query(self, num_examples, vector_index: type[VectorStoreBase]): - self.operators.append( - GremlinExampleIndexQuery(vector_index, self.embedding, num_examples) - ) + self.operators.append(GremlinExampleIndexQuery(vector_index, self.embedding, num_examples)) return self def gremlin_generate_synthesize( self, schema, gremlin_prompt: Optional[str] = None, vertices: Optional[List[str]] = None ): - self.operators.append( - GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt) - ) + self.operators.append(GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt)) return self def print_result(self): 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 aa886c0af..52626b72b 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 @@ -50,9 +50,7 @@ def run(self, data: dict) -> Dict[str, Any]: if not schema: # TODO: ensure the function works correctly (update the logic later) self.schema_free_mode(data.get("triples", [])) - log.warning( - "Using schema_free mode, could try schema_define mode for better effect!" - ) + log.warning("Using schema_free mode, could try schema_define mode for better effect!") else: self.init_schema_if_need(schema) self.load_into_graph(vertices, edges, schema) @@ -68,9 +66,7 @@ def _set_default_property(self, key, input_properties, property_label_map): # list or set default_value = [] input_properties[key] = default_value - log.warning( - "Property '%s' missing in vertex, set to '%s' for now", key, default_value - ) + log.warning("Property '%s' missing in vertex, set to '%s' for now", key, default_value) def _handle_graph_creation(self, func, *args, **kwargs): try: @@ -82,17 +78,11 @@ def _handle_graph_creation(self, func, *args, **kwargs): log.error("Error on creating: %s, %s", args, e) return None - def load_into_graph( - self, vertices, edges, schema - ): # pylint: disable=too-many-statements + def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many-statements # pylint: disable=R0912 (too-many-branches) - vertex_label_map = { - v_label["name"]: v_label for v_label in schema["vertexlabels"] - } + 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"] - } + property_label_map = {p_label["name"]: p_label for p_label in schema["propertykeys"]} for vertex in vertices: input_label = vertex["label"] @@ -108,9 +98,7 @@ def load_into_graph( vertex_label = vertex_label_map[input_label] primary_keys = vertex_label["primary_keys"] nullable_keys = vertex_label.get("nullable_keys", []) - non_null_keys = [ - key for key in vertex_label["properties"] if key not in nullable_keys - ] + non_null_keys = [key for key in vertex_label["properties"] if key not in nullable_keys] has_problem = False # 2. Handle primary-keys mode vertex @@ -142,9 +130,7 @@ def load_into_graph( # 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 - ) + 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(): @@ -181,9 +167,7 @@ def load_into_graph( 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 - ) + self._handle_graph_creation(self.client.graph().addEdge, label, start, end, properties) def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -207,20 +191,18 @@ def init_schema_if_need(self, schema: dict): source_vertex_label = edge["source_label"] target_vertex_label = edge["target_label"] properties = edge["properties"] - self.schema.edgeLabel(edge_label).sourceLabel( - source_vertex_label - ).targetLabel(target_vertex_label).properties(*properties).nullableKeys( - *properties - ).ifNotExist().create() + self.schema.edgeLabel(edge_label).sourceLabel(source_vertex_label).targetLabel( + target_vertex_label + ).properties(*properties).nullableKeys(*properties).ifNotExist().create() def schema_free_mode(self, 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("name").ifNotExist().create() + self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel("vertex").properties( + "name" + ).ifNotExist().create() self.schema.indexLabel("vertexByName").onV("vertex").by( "name" @@ -280,9 +262,7 @@ def _set_property_data_type(self, property_key, data_type): log.warning("UUID type is not supported, use text instead") property_key.asText() else: - log.error( - "Unknown data type %s for property_key %s", data_type, property_key - ) + log.error("Unknown data type %s for property_key %s", data_type, property_key) def _set_property_cardinality(self, property_key, cardinality): if cardinality == PropertyCardinality.SINGLE: @@ -292,13 +272,9 @@ def _set_property_cardinality(self, property_key, cardinality): elif cardinality == PropertyCardinality.SET: property_key.valueSet() else: - log.error( - "Unknown cardinality %s for property_key %s", cardinality, property_key - ) + 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: + def _check_property_data_type(self, data_type: str, cardinality: str, value) -> bool: if cardinality in ( PropertyCardinality.LIST.value, PropertyCardinality.SET.value, @@ -328,9 +304,7 @@ def _check_single_data_type(self, data_type: str, value) -> bool: if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below - if ( - data_type == PropertyDataType.DATE.value - ): # the format should be "yyyy-MM-dd" + if data_type == PropertyDataType.DATE.value: # the format should be "yyyy-MM-dd" import re return isinstance(value, str) and re.match(r"^\d{4}-\d{2}-\d{2}$", value) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py index 73c9530df..4c4c167c4 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py @@ -47,7 +47,5 @@ def res = [:]; result = self.graph.gremlin().exec(groovy_code)["data"] if isinstance(result, list) and len(result) > 0: - graph_summary.update( - {key: result[i].get(key) for i, key in enumerate(keys)} - ) + graph_summary.update({key: result[i].get(key) for i, key in enumerate(keys)}) return graph_summary diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py index 52399d99e..d9542cd39 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py @@ -132,9 +132,7 @@ def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: query_embedding = context.get("query_embedding") self._gremlin_generator.clear() - self._gremlin_generator.example_index_query( - num_examples=self._num_gremlin_generate_example - ) + self._gremlin_generator.example_index_query(num_examples=self._num_gremlin_generate_example) gremlin_response = self._gremlin_generator.gremlin_generate_synthesize( context["simple_schema"], vertices=vertices, gremlin_prompt=self._gremlin_prompt ).run(query=query, query_embedding=query_embedding) @@ -148,9 +146,7 @@ def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: result = self._client.gremlin().exec(gremlin=gremlin)["data"] if result == [None]: result = [] - context["graph_result"] = [ - json.dumps(item, ensure_ascii=False) for item in result - ] + context["graph_result"] = [json.dumps(item, ensure_ascii=False) for item in result] if context["graph_result"]: context["graph_result_flag"] = 1 context["graph_context_head"] = ( @@ -228,9 +224,7 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: "Unable to find vid, downgraded to property query, please confirm if it meets expectation." ) - paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)[ - "data" - ] + paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)["data"] graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( self._format_graph_query_result(query_paths=paths) ) @@ -352,18 +346,14 @@ def _process_vertex( use_id_to_match: bool, v_cache: Set[str], ) -> Tuple[str, int, int]: - matched_str = ( - item["id"] if use_id_to_match else item["props"][self._prop_to_match] - ) + matched_str = item["id"] if use_id_to_match else item["props"][self._prop_to_match] if matched_str in node_cache: flat_rel = flat_rel[:-prior_edge_str_len] return flat_rel, prior_edge_str_len, depth node_cache.add(matched_str) props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'v')}" - for k, v in item["props"].items() - if v + f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v ) # TODO: we may remove label id or replace with label name @@ -388,9 +378,7 @@ def _process_edge( e_cache: Set[Tuple[str, str, str]], ) -> Tuple[str, int]: props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'e')}" - for k, v in item["props"].items() - if v + f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v ) props_str = f"{{{props_str}}}" if props_str else "" prev_matched_str = ( @@ -407,9 +395,7 @@ def _process_edge( edge_label = item["label"] edge_str = ( - f"--[{edge_label}]-->" - if item["outV"] == prev_matched_str - else f"<--[{edge_label}]--" + f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" ) path_str += edge_str prior_edge_str_len = len(edge_str) @@ -427,20 +413,14 @@ def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: schema = self._get_graph_schema() vertex_props_str, edge_props_str = schema.split("\n")[:2] # TODO: rename to vertex (also need update in the schema) - vertex_props_str = ( - vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") - ) - edge_props_str = ( - edge_props_str[len("Edge properties: ") :].strip("[").strip("]") - ) + vertex_props_str = vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") + edge_props_str = edge_props_str[len("Edge properties: ") :].strip("[").strip("]") vertex_labels = self._extract_label_names(vertex_props_str) edge_labels = self._extract_label_names(edge_props_str) return vertex_labels, edge_labels @staticmethod - def _extract_label_names( - source: str, head: str = "name: ", tail: str = ", " - ) -> List[str]: + def _extract_label_names(source: str, head: str = "name: ", tail: str = ", ") -> List[str]: result = [] for s in source.split(head): end = s.find(tail) @@ -466,9 +446,7 @@ def _get_graph_schema(self, refresh: bool = False) -> str: log.debug("Link(Relation): %s", relationships) return self._schema - def _limit_property_query( - self, value: Optional[str], item_type: str - ) -> Optional[str]: + def _limit_property_query(self, value: Optional[str], item_type: str) -> Optional[str]: # NOTE: we skip the filter for list/set type (e.g., list of string, add it if needed) if not self._limit_property or not isinstance(value, str): return value 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 e9bccc2f4..2f0643a77 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 @@ -40,9 +40,7 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: new_vertex = { - key: vertex[key] - for key in ["id", "name", "properties"] - if key in vertex + key: vertex[key] for key in ["id", "name", "properties"] if key in vertex } mini_schema["vertexlabels"].append(new_vertex) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index d41f99a37..4f288a9de 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -62,14 +62,10 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: # !: We have assumed that self.example is not empty queries = [example["query"] for example in self.examples] # TODO: refactor function chain async to avoid blocking - examples_embedding = asyncio.run( - get_embeddings_parallel(self.embedding, queries) - ) + examples_embedding = asyncio.run(get_embeddings_parallel(self.embedding, queries)) embed_dim = len(examples_embedding[0]) if len(self.examples) > 0: - vector_index = self.vector_index.from_name( - embed_dim, self.vector_index_name - ) + vector_index = self.vector_index.from_name(embed_dim, self.vector_index_name) vector_index.add(examples_embedding, self.examples) <<<<<<< HEAD vector_index.to_index_file(self.index_dir, self.filename_prefix) 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 3bd838168..4c46b8906 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 @@ -49,9 +49,7 @@ async def get_embeddings_with_semaphore(vid_list: list[str]) -> Any: None, self.embedding.get_texts_embeddings, vid_list ) - vid_batches = [ - vids[i : i + batch_size] for i in range(0, len(vids), batch_size) - ] + vid_batches = [vids[i : i + batch_size] for i in range(0, len(vids), batch_size)] tasks = [get_embeddings_with_semaphore(batch) for batch in vid_batches] embeddings = [] @@ -64,26 +62,18 @@ async def get_embeddings_with_semaphore(vid_list: list[str]) -> Any: def run(self, context: Dict[str, Any]) -> Dict[str, Any]: vertexlabels = self.sm.schema.getSchema()["vertexlabels"] - all_pk_flag = all( - data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels - ) + all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) past_vids = self.vid_index.get_all_properties() # only support Faiss # TODO: We should build vid vector index separately, especially when the vertices may be very large - present_vids = context[ - "vertices" - ] # Warning: data truncated by fetch_graph_data.py + present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py removed_vids = set(past_vids) - set(present_vids) removed_num = self.vid_index.remove(removed_vids) added_vids = list(set(present_vids) - set(past_vids)) if added_vids: - vids_to_process = ( - self._extract_names(added_vids) if all_pk_flag else added_vids - ) - added_embeddings = asyncio.run( - self._get_embeddings_parallel(vids_to_process) - ) + vids_to_process = self._extract_names(added_vids) if all_pk_flag else added_vids + 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") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index 8d1f3394f..b055cf62e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -59,23 +59,19 @@ def __init__( self.embedding.get_embedding_dim(), "gremlin_examples" ) - def _get_match_result( - self, context: Dict[str, Any], query: str - ) -> List[Dict[str, Any]]: + def _get_match_result(self, context: Dict[str, Any], query: str) -> List[Dict[str, Any]]: if self.num_examples <= 0: return [] query_embedding = context.get("query_embedding") if not isinstance(query_embedding, list): query_embedding = self.embedding.get_texts_embeddings([query])[0] - return self.vector_index.search( - query_embedding, self.num_examples, dis_threshold=1.8 - ) + return self.vector_index.search(query_embedding, self.num_examples, dis_threshold=1.8) def _build_default_example_index(self): - properties = pd.read_csv( - os.path.join(resource_path, "demo", "text2gremlin.csv") - ).to_dict(orient="records") + properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict( + orient="records" + ) from concurrent.futures import ThreadPoolExecutor # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py index 923068462..49e712d01 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py @@ -39,9 +39,7 @@ def __init__( topk_per_keyword: int = huge_settings.topk_per_keyword, vector_dis_threshold: float = huge_settings.vector_dis_threshold, ): - self.index_dir = str( - os.path.join(resource_path, huge_settings.graph_name, "graph_vids") - ) + self.index_dir = str(os.path.join(resource_path, huge_settings.graph_name, "graph_vids")) self.vector_index = vector_index.from_name( embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids" ) @@ -67,9 +65,7 @@ def _exact_match_vids(self, keywords: List[str]) -> Tuple[List[str], List[str]]: possible_vids.update([f"{i + 1}:{keyword}" for keyword in keywords]) vids_str = ",".join([f"'{vid}'" for vid in possible_vids]) - resp = self._client.gremlin().exec( - SemanticIdQuery.ID_QUERY_TEMPL.format(vids_str=vids_str) - ) + resp = self._client.gremlin().exec(SemanticIdQuery.ID_QUERY_TEMPL.format(vids_str=vids_str)) searched_vids = [v["id"] for v in resp["data"]] unsearched_keywords = set(keywords) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py index cb7d1f1ae..0b396b77d 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py @@ -51,9 +51,7 @@ def __init__( self.graph = graph self.result = None - def import_schema( - self, from_hugegraph=None, from_extraction=None, from_user_defined=None - ): + def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): if from_hugegraph: self.operators.append(SchemaManager(from_hugegraph)) elif from_user_defined: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py index 2447c2e06..2140abca2 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py @@ -62,10 +62,8 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = ( - f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( - "\n" - ) + context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( + "\n" ) final_prompt = self._prompt_template.format( @@ -92,12 +90,8 @@ def init_llm(self, context): if self._question is None: self._question = context.get("query") or None assert self._question is not None, "No question for synthesizing." - context_head_str = ( - context.get("synthesize_context_head") or self._context_head or "" - ) - context_tail_str = ( - context.get("synthesize_context_tail") or self._context_tail or "" - ) + context_head_str = context.get("synthesize_context_head") or self._context_head or "" + context_tail_str = context.get("synthesize_context_tail") or self._context_tail or "" return context_head_str, context_tail_str def handle_vector_graph(self, context): @@ -121,16 +115,12 @@ def handle_vector_graph(self, context): log.warning(graph_result_context) return graph_result_context, vector_result_context - async def run_streaming( - self, context: Dict[str, Any] - ) -> AsyncGenerator[Dict[str, Any], None]: + async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[str, Any], None]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = ( - f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( - "\n" - ) + context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( + "\n" ) final_prompt = self._prompt_template.format( @@ -163,9 +153,7 @@ async def async_generate( async_tasks = {} if self._raw_answer: final_prompt = self._question - async_tasks["raw_task"] = asyncio.create_task( - self._llm.agenerate(prompt=final_prompt) - ) + async_tasks["raw_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) if self._vector_only_answer: context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip( "\n" @@ -178,10 +166,8 @@ async def async_generate( self._llm.agenerate(prompt=final_prompt) ) if self._graph_only_answer: - context_str = ( - f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( - "\n" - ) + context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( + "\n" ) final_prompt = self._prompt_template.format( @@ -194,11 +180,7 @@ async def async_generate( context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = ( - f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip( - "\n" - ) - ) + context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -267,10 +249,8 @@ async def async_streaming_generate( ) auto_id += 1 if self._graph_only_answer: - context_str = ( - f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( - "\n" - ) + context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( + "\n" ) final_prompt = self._prompt_template.format( @@ -286,11 +266,7 @@ async def async_streaming_generate( context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = ( - f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip( - "\n" - ) - ) + context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -316,9 +292,7 @@ async def async_streaming_generate( async_tasks = [asyncio.create_task(anext(gen)) for gen in async_generators] while True: - done, _ = await asyncio.wait( - async_tasks, return_when=asyncio.FIRST_COMPLETED - ) + done, _ = await asyncio.wait(async_tasks, return_when=asyncio.FIRST_COMPLETED) stop_task_num = 0 for task in done: try: @@ -332,9 +306,7 @@ async def async_streaming_generate( break yield context - async def __llm_generate_with_meta_info( - self, task_id: int, target_key: str, prompt: str - ): + async def __llm_generate_with_meta_info(self, task_id: int, target_key: str, prompt: str): # FIXME: Expected type 'AsyncIterable', got 'Coroutine[Any, Any, AsyncGenerator[str, None]]' instead async for token in self._llm.agenerate_streaming(prompt=prompt): yield task_id, target_key, token diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index 2c0244d57..fd4583263 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -48,9 +48,7 @@ def _extract_response(self, response: str, label: str = "gremlin") -> str: return match.group(1).strip() return response.strip() - def _format_examples( - self, examples: Optional[List[Dict[str, str]]] - ) -> Optional[str]: + def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional[str]: if not examples: return None example_strings = [] @@ -88,9 +86,7 @@ def _format_properties(self, properties: Optional[List[tuple]]) -> Optional[str] async def async_generate(self, context: Dict[str, Any]): async_tasks = {} query = context.get("query") - raw_example = [ - {"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"} - ] + raw_example = [{"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"}] raw_prompt = self.gremlin_prompt.format( query=query, schema=self.schema, @@ -98,9 +94,7 @@ async def async_generate(self, context: Dict[str, Any]): vertices=self._format_vertices(vertices=self.vertices), properties=self._format_properties(properties=None), ) - async_tasks["raw_answer"] = asyncio.create_task( - self.llm.agenerate(prompt=raw_prompt) - ) + async_tasks["raw_answer"] = asyncio.create_task(self.llm.agenerate(prompt=raw_prompt)) examples = context.get("match_result") init_prompt = self.gremlin_prompt.format( @@ -130,9 +124,7 @@ async def async_generate(self, context: Dict[str, Any]): def sync_generate(self, context: Dict[str, Any]): query = context.get("query") - raw_example = [ - {"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"} - ] + raw_example = [{"query": "who is peter", "gremlin": "g.V().has('name', 'peter')"}] raw_prompt = self.gremlin_prompt.format( query=query, schema=self.schema, 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 609c5f22f..a9ac4c050 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 @@ -197,12 +197,8 @@ def valid(self, element_id: str, max_length: int = 256) -> bool: return True def _filter_long_id(self, graph) -> Dict[str, List[Any]]: - graph["vertices"] = [ - vertex for vertex in graph["vertices"] if self.valid(vertex["id"]) - ] + graph["vertices"] = [vertex for vertex in graph["vertices"] if self.valid(vertex["id"])] graph["edges"] = [ - edge - for edge in graph["edges"] - if self.valid(edge["start"]) and self.valid(edge["end"]) + edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"]) ] return graph diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index 420fc9776..0fa1c0f04 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -113,8 +113,6 @@ def _extract_keywords_from_response( sub_tokens = re.findall(r"\w+", token) if len(sub_tokens) > 1: results.update( - w - for w in sub_tokens - if w not in NLTKHelper().stopwords(lang=self._language) + w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language) ) return results diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py index a45812393..058d1bce9 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py @@ -31,9 +31,7 @@ def __init__(self, llm: BaseLLM): def _load_few_shot_example(self, example_name: str) -> Dict[str, Any]: """Loads and finds the specified few-shot example from the unified JSON file.""" - examples_path = os.path.join( - resource_path, "prompt_examples", "prompt_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") if not os.path.exists(examples_path): raise FileNotFoundError(f"Examples file not found: {examples_path}") with open(examples_path, "r", encoding="utf-8") as f: @@ -41,9 +39,7 @@ def _load_few_shot_example(self, example_name: str) -> Dict[str, Any]: for example in all_examples: if example.get("name") == example_name: return example - raise ValueError( - f"Example with name '{example_name}' not found in prompt_examples.json" - ) + raise ValueError(f"Example with name '{example_name}' not found in prompt_examples.json") def run(self, context: Dict[str, Any]) -> Dict[str, Any]: """Executes the core logic of prompt generation.""" @@ -52,9 +48,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: example_name = context.get("example_name") if not all([source_text, scenario, example_name]): - raise ValueError( - "Missing required context: source_text, scenario, or example_name." - ) + raise ValueError("Missing required context: source_text, scenario, or example_name.") few_shot_example = self._load_few_shot_example(example_name) meta_prompt = prompt_tpl.generate_extract_prompt_template.format( 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 3ba178c88..d517db8b9 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 @@ -67,9 +67,9 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: item_type = item["type"] if item_type == "vertex": label = item["label"] - non_nullable_keys = set( - properties_map[item_type][label]["properties"] - ).difference(set(properties_map[item_type][label]["nullable_keys"])) + non_nullable_keys = set(properties_map[item_type][label]["properties"]).difference( + set(properties_map[item_type][label]["nullable_keys"]) + ) for key in non_nullable_keys: if key not in item["properties"]: item["properties"][key] = "NULL" @@ -82,9 +82,7 @@ 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) -> None: self.llm = llm self.example_prompt = example_prompt self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name @@ -150,9 +148,7 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: and "vertices" in property_graph and "edges" in property_graph ): - log.critical( - "Invalid property graph format; expecting 'vertices' and 'edges'." - ) + log.critical("Invalid property graph format; expecting 'vertices' and 'edges'.") return items # Create sets for valid vertex and edge labels based on the schema @@ -162,9 +158,7 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: def process_items(item_list, valid_labels, item_type): for item in item_list: if not isinstance(item, dict): - log.warning( - "Invalid property graph item type '%s'.", type(item) - ) + log.warning("Invalid property graph item type '%s'.", type(item)) continue if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): log.warning("Invalid item keys '%s'.", item.keys()) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py index ae445c206..1e33514ca 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py @@ -117,13 +117,9 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: raise ValueError("Context must be a dictionary") if "raw_texts" not in context or not isinstance(context["raw_texts"], list): raise ValueError("'raw_texts' must be a list[str]") - if "query_examples" not in context or not isinstance( - context["query_examples"], list - ): + if "query_examples" not in context or not isinstance(context["query_examples"], list): raise ValueError("'query_examples' must be a list[str]") - if "few_shot_schema" not in context or not isinstance( - context["few_shot_schema"], dict - ): + if "few_shot_schema" not in context or not isinstance(context["few_shot_schema"], dict): raise ValueError("'few_shot_schema' must be a dict") raw_texts = context["raw_texts"] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py index 98cc97ccf..6beeb0291 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py @@ -104,9 +104,7 @@ def nodes_schemas_text_to_list_of_dict(nodes_schemas): properties = json.loads(properties) except json.decoder.JSONDecodeError: properties = {} - result.append( - {"label": label, "primary_key": primary_key, "properties": properties} - ) + result.append({"label": label, "primary_key": primary_key, "properties": properties}) return result @@ -118,9 +116,7 @@ def relationships_schemas_text_to_list_of_dict(relationships_schemas): continue start = relationships_schema_list[0].strip().replace('"', "") end = relationships_schema_list[2].strip().replace('"', "") - relationships_schema_type = ( - relationships_schema_list[1].strip().replace('"', "") - ) + relationships_schema_type = relationships_schema_list[1].strip().replace('"', "") properties = re.search(JSON_REGEX, relationships_schema) if properties is None: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py index b5232d268..2914c4b28 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py @@ -23,9 +23,7 @@ from hugegraph_llm.utils.log import log -def log_elapsed_time( - start_time: float, func: Callable, args: tuple, msg: Optional[str] -): +def log_elapsed_time(start_time: float, func: Callable, args: tuple, msg: Optional[str]): elapse_time = time.perf_counter() - start_time unit = "s" if elapse_time < 1: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py index ace3d4b6a..b2f485cea 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py @@ -32,9 +32,7 @@ async def _get_batch_with_progress( return result -async def get_embeddings_parallel( - embedding: BaseEmbedding, vids: list[str] -) -> list[Any]: +async def get_embeddings_parallel(embedding: BaseEmbedding, vids: list[str]) -> list[Any]: """Get embeddings for texts in parallel. This function processes text embeddings asynchronously in parallel, using batching and semaphore @@ -62,9 +60,7 @@ async def get_embeddings_parallel( embeddings = [] with tqdm(total=len(vid_batches)) as pbar: # Create tasks for each batch with progress bar updates - tasks = [ - _get_batch_with_progress(embedding, batch, pbar) for batch in vid_batches - ] + tasks = [_get_batch_with_progress(embedding, batch, pbar) for batch in vid_batches] # Use asyncio.gather() to preserve order batch_results = await asyncio.gather(*tasks) @@ -78,9 +74,7 @@ async def get_embeddings_parallel( def get_filename_prefix(embedding_type: str = None, model_name: str = None) -> str: """Generate filename based on model name.""" - if not ( - model_name and model_name.strip() and embedding_type and embedding_type.strip() - ): + if not (model_name and model_name.strip() and embedding_type and embedding_type.strip()): return "" # Sanitize model_name to prevent path traversal or invalid filename chars safe_embedding_type = embedding_type.replace("/", "_").replace("\\", "_").strip() diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index e7b213c61..2f48af332 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -33,9 +33,7 @@ def get_graph_index_info(): - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) graph_summary_info = builder.fetch_graph_data().run() vector_index = get_vector_index_class(index_settings.cur_vector_index) vector_index_entity = vector_index.from_name( @@ -90,18 +88,14 @@ def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: def extract_graph(input_file, input_text, schema, example_prompt) -> str: texts = read_documents(input_file, input_text) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) if not schema: return "ERROR: please input with correct schema/format." error_message = parse_schema(schema, builder) if error_message: return error_message - builder.chunk_split(texts, "document", "zh").extract_info( - example_prompt, "property_graph" - ) + builder.chunk_split(texts, "document", "zh").extract_info(example_prompt, "property_graph") try: context = builder.run() @@ -128,9 +122,7 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: def update_vid_embedding(): vector_index = get_vector_index_class(index_settings.cur_vector_index) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) builder.fetch_graph_data().build_vertex_id_semantic_index(vector_index) log.debug("Operators: %s", builder.operators) try: @@ -147,9 +139,7 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: data_json = json.loads(data.strip()) log.debug("Import graph data: %s", data) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) if schema: error_message = parse_schema(schema, builder) if error_message: @@ -195,9 +185,7 @@ def build_schema(input_text, query_example, few_shot): except json.JSONDecodeError as e: raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) try: schema = builder.build_schema().run(context) except Exception as e: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index 082d6bb8c..6da7a6567 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -29,9 +29,7 @@ MAX_BACKUP_DIRS = 7 MAX_VERTICES = 100000 MAX_EDGES = 200000 -BACKUP_DIR = str( - os.path.join(resource_path, "backup-graph-data-4020", huge_settings.graph_name) -) +BACKUP_DIR = str(os.path.join(resource_path, "backup-graph-data-4020", huge_settings.graph_name)) def run_gremlin_query(query, fmt=True): @@ -58,33 +56,21 @@ def init_hg_test_data(): schema.vertexLabel("Person").properties( "name", "birthDate" ).useCustomizeStringId().ifNotExist().create() - schema.vertexLabel("Movie").properties( - "name" - ).useCustomizeStringId().ifNotExist().create() - schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel( - "Movie" - ).ifNotExist().create() + schema.vertexLabel("Movie").properties("name").useCustomizeStringId().ifNotExist().create() + schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() - schema.indexLabel("PersonByName").onV("Person").by( - "name" - ).secondary().ifNotExist().create() - schema.indexLabel("MovieByName").onV("Movie").by( - "name" - ).secondary().ifNotExist().create() + schema.indexLabel("PersonByName").onV("Person").by("name").secondary().ifNotExist().create() + schema.indexLabel("MovieByName").onV("Movie").by("name").secondary().ifNotExist().create() graph = client.graph() - graph.addVertex( - "Person", {"name": "Al Pacino", "birthDate": "1940-04-25"}, id="Al Pacino" - ) + graph.addVertex("Person", {"name": "Al Pacino", "birthDate": "1940-04-25"}, id="Al Pacino") graph.addVertex( "Person", {"name": "Robert De Niro", "birthDate": "1943-08-17"}, id="Robert De Niro", ) graph.addVertex("Movie", {"name": "The Godfather"}, id="The Godfather") - graph.addVertex( - "Movie", {"name": "The Godfather Part II"}, id="The Godfather Part II" - ) + graph.addVertex("Movie", {"name": "The Godfather Part II"}, id="The Godfather Part II") graph.addVertex( "Movie", {"name": "The Godfather Coda The Death of Michael Corleone"}, @@ -93,9 +79,7 @@ def init_hg_test_data(): graph.addEdge("ActedIn", "Al Pacino", "The Godfather", {}) graph.addEdge("ActedIn", "Al Pacino", "The Godfather Part II", {}) - graph.addEdge( - "ActedIn", "Al Pacino", "The Godfather Coda The Death of Michael Corleone", {} - ) + graph.addEdge("ActedIn", "Al Pacino", "The Godfather Coda The Death of Michael Corleone", {}) graph.addEdge("ActedIn", "Robert De Niro", "The Godfather Part II", {}) schema.getSchema() graph.close() @@ -134,9 +118,7 @@ def backup_data(): } vertexlabels = client.schema().getSchema()["vertexlabels"] - all_pk_flag = all( - data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels - ) + all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) for filename, query in files.items(): write_backup_file(client, backup_subdir, filename, query, all_pk_flag) @@ -216,9 +198,7 @@ def manage_backup_retention(): # TODO: In the path demo/rag_demo/configs_block.py, # there is a function test_api_connection that is similar to this function, # but it is not straightforward to reuse -def check_graph_db_connection( - url: str, name: str, user: str, pwd: str, graph_space: str -) -> bool: +def check_graph_db_connection(url: str, name: str, user: str, pwd: str, graph_space: str) -> bool: try: if graph_space and graph_space.strip(): test_url = f"{url}/graphspaces/{graph_space}/graphs/{name}/schema" diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index df48ea1e9..4f56db636 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -51,9 +51,7 @@ def read_documents(input_file, input_text): texts.append(text) elif full_path.endswith(".pdf"): # TODO: support PDF file - raise gr.Error( - "PDF will be supported later! Try to upload text/docx now" - ) + raise gr.Error("PDF will be supported later! Try to upload text/docx now") else: raise gr.Error("Please input txt or docx file.") else: @@ -89,14 +87,8 @@ def build_vector_index(input_file, input_text): if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) - context = ( - builder.chunk_split(texts, "paragraph", "zh") - .build_vector_index(vector_index) - .run() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + context = builder.chunk_split(texts, "paragraph", "zh").build_vector_index(vector_index).run() return json.dumps(context, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/tests/models/embeddings/test_ollama_embedding.py b/hugegraph-llm/src/tests/models/embeddings/test_ollama_embedding.py index 3abccc268..a7a9d044c 100644 --- a/hugegraph-llm/src/tests/models/embeddings/test_ollama_embedding.py +++ b/hugegraph-llm/src/tests/models/embeddings/test_ollama_embedding.py @@ -32,7 +32,5 @@ def test_get_cosine_similarity(self): ollama_embedding = OllamaEmbedding(model_name="quentinz/bge-large-zh-v1.5") embedding1 = ollama_embedding.get_text_embedding("hello world") embedding2 = ollama_embedding.get_text_embedding("bye world") - similarity = OllamaEmbedding.similarity( - embedding1, embedding2, SimilarityMode.DEFAULT - ) + similarity = OllamaEmbedding.similarity(embedding1, embedding2, SimilarityMode.DEFAULT) print(similarity) diff --git a/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py b/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py index ad975a37e..317d02879 100644 --- a/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py +++ b/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py @@ -26,9 +26,7 @@ def setUp(self): def test_schema_check_with_valid_input(self): data = { - "vertexlabels": [ - {"name": "person", "properties": ["name", "age", "occupation"]} - ], + "vertexlabels": [{"name": "person", "properties": ["name", "age", "occupation"]}], "edgelabels": [ { "name": "knows", diff --git a/hugegraph-llm/src/tests/operators/llm_op/test_disambiguate_data.py b/hugegraph-llm/src/tests/operators/llm_op/test_disambiguate_data.py index e120a1b00..04ee42142 100644 --- a/hugegraph-llm/src/tests/operators/llm_op/test_disambiguate_data.py +++ b/hugegraph-llm/src/tests/operators/llm_op/test_disambiguate_data.py @@ -57,9 +57,7 @@ def setUp(self): "properties": {"name": "www.bob.com", "url": "www.bob.com"}, }, ], - "edges": [ - {"start": "Alice", "end": "Bob", "type": "roommate", "properties": {}} - ], + "edges": [{"start": "Alice", "end": "Bob", "type": "roommate", "properties": {}}], "schema": { "vertices": [ {"vertex_label": "person", "properties": ["name", "age", "occupation"]}, From b78b051522b511fc5222148805f67397cbc33f32 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Thu, 9 Oct 2025 15:27:50 +0800 Subject: [PATCH 30/71] fix: resolve leftover conflict markers and deps in hugegraph-llm/pyproject.toml --- hugegraph-llm/pyproject.toml | 9 --------- 1 file changed, 9 deletions(-) diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index 7724f7444..4b266cd20 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -41,7 +41,6 @@ dependencies = [ "pandas", "pydantic", "tqdm", - "tqdm", # LLM specific dependencies "openai", @@ -69,15 +68,7 @@ vectordb = [ "pymilvus==2.5.9", "qdrant-client==1.14.2", ] -======= - - # Vector database dependencies - "pymilvus==2.5.9", - "qdrant-client==1.14.2", - - ] ->>>>>>> a255aed (fix cycle import & add docs) [project.urls] homepage = "https://hugegraph.apache.org/" repository = "https://github.com/apache/incubator-hugegraph-ai" From 085e372e8997294e825fdbf033c7cf969d1c97da Mon Sep 17 00:00:00 2001 From: lingxiao Date: Thu, 9 Oct 2025 15:30:02 +0800 Subject: [PATCH 31/71] chore: clean remaining conflict markers by preferring PR-side chunks --- .../demo/rag_demo/admin_block.py | 8 ---- .../demo/rag_demo/other_block.py | 6 --- .../src/hugegraph_llm/document/chunk_split.py | 18 ------- .../hugegraph_llm/models/embeddings/base.py | 24 ---------- .../models/embeddings/init_embedding.py | 47 ------------------- .../models/embeddings/litellm.py | 9 ---- .../hugegraph_llm/models/embeddings/ollama.py | 34 -------------- .../hugegraph_llm/models/embeddings/openai.py | 7 --- .../src/hugegraph_llm/models/llms/ollama.py | 8 ---- .../hugegraph_llm/operators/graph_rag_task.py | 9 ---- .../index_op/build_gremlin_example_index.py | 32 ------------- .../index_op/gremlin_example_index_query.py | 6 --- .../operators/index_op/vector_index_query.py | 20 -------- .../operators/llm_op/disambiguate_data.py | 9 ---- .../operators/llm_op/gremlin_generate.py | 9 ---- .../operators/llm_op/keyword_extract.py | 8 ---- .../llm_op/property_graph_extract.py | 6 --- .../operators/llm_op/schema_build.py | 7 --- .../hugegraph_llm/utils/hugegraph_utils.py | 8 ---- hugegraph-llm/src/hugegraph_llm/utils/log.py | 8 ---- .../tests/models/llms/test_ollama_client.py | 6 --- .../src/pyhugegraph/api/graph.py | 14 ------ .../pyhugegraph/example/hugegraph_example.py | 8 ---- .../src/pyhugegraph/utils/util.py | 8 ---- 24 files changed, 319 deletions(-) mode change 100755 => 100644 hugegraph-llm/src/hugegraph_llm/utils/log.py diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 1a157d387..00da2a973 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -108,15 +108,7 @@ def create_admin_block(): ) # Error message box, initially hidden -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - error_message = gr.Textbox( - label="", visible=False, interactive=False, elem_classes="error-message" - ) -======= error_message = gr.Textbox(label="", visible=False, interactive=False, elem_classes="error-message") ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ======= error_message = gr.Textbox(label="", visible=False, interactive=False, elem_classes="error-message") >>>>>>> 8e0bf08 (chore: mark vectordb optional) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py index 82650b907..efa77602e 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py @@ -58,13 +58,7 @@ def create_other_block(): async def lifespan(app: FastAPI): # pylint: disable=W0621 log.info("Starting background scheduler...") scheduler = AsyncIOScheduler() -<<<<<<< HEAD - scheduler.add_job( - backup_data, trigger=CronTrigger(hour=1, minute=0), id="daily_backup", replace_existing=True - ) -======= scheduler.add_job(backup_data, trigger=CronTrigger(hour=1, minute=0), id="daily_backup", replace_existing=True) ->>>>>>> 87ee5d3 (style: format code with black line-length 120) scheduler.start() log.info("Starting vid embedding update task...") diff --git a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py index 11369c06a..c9ece77bb 100644 --- a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py +++ b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py @@ -22,13 +22,7 @@ class ChunkSplitter: def __init__( -<<<<<<< HEAD - self, - split_type: Literal["paragraph", "sentence"] = "paragraph", - language: Literal["zh", "en"] = "zh", -======= self, split_type: Literal["paragraph", "sentence"] = "paragraph", language: Literal["zh", "en"] = "zh" ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ): if language == "zh": separators = ["\n\n", "\n", "。", "，", ""] @@ -37,21 +31,9 @@ def __init__( else: raise ValueError("Argument `language` must be zh or en!") if split_type == "paragraph": -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - self.text_splitter = RecursiveCharacterTextSplitter( - chunk_size=500, chunk_overlap=30, separators=separators - ) - elif split_type == "sentence": - self.text_splitter = RecursiveCharacterTextSplitter( - chunk_size=50, chunk_overlap=0, separators=separators - ) -======= self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=30, separators=separators) elif split_type == "sentence": self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=50, chunk_overlap=0, separators=separators) ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ======= self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=30, separators=separators) elif split_type == "sentence": diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py index f895314f3..62874319c 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py @@ -61,15 +61,12 @@ def get_text_embedding(self, text: str) -> List[float]: """Comment""" @abstractmethod -<<<<<<< HEAD -======= def get_embedding_dim( self, ) -> int: """Comment""" @abstractmethod ->>>>>>> 38dce0b (feat(llm): vector db finished) def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch. @@ -90,29 +87,8 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """ @abstractmethod -<<<<<<< HEAD - async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts in a single batch asynchronously. - - This method should efficiently process multiple texts at once by leveraging - the embedding model's batching capabilities, which is typically more efficient - than processing texts individually. - - Parameters - ---------- - texts : List[str] - A list of text strings to be embedded. - - Returns - ------- - List[List[float]] - A list of embedding vectors, where each vector is a list of floats. - The order of embeddings should match the order of input texts. - """ -======= async def async_get_text_embedding(self, text: str) -> List[float]: """Comment""" ->>>>>>> 38dce0b (feat(llm): vector db finished) @staticmethod def similarity( diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index d96840911..143d54229 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -17,46 +17,10 @@ from hugegraph_llm.config import llm_settings -<<<<<<< HEAD -from hugegraph_llm.config import LLMConfig -from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding -from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding -from hugegraph_llm.models.embeddings.openai import OpenAIEmbedding - -model_map = { - "openai": llm_settings.openai_embedding_model, - "ollama/local": llm_settings.ollama_embedding_model, - "litellm": llm_settings.litellm_embedding_model, -} - - -def get_embedding(llm_settings: LLMConfig): - if llm_settings.embedding_type == "openai": - return OpenAIEmbedding( - model_name=llm_settings.openai_embedding_model, - api_key=llm_settings.openai_embedding_api_key, - api_base=llm_settings.openai_embedding_api_base, - ) - if llm_settings.embedding_type == "ollama/local": - return OllamaEmbedding( - model_name=llm_settings.ollama_embedding_model, - host=llm_settings.ollama_embedding_host, - port=llm_settings.ollama_embedding_port, - ) - if llm_settings.embedding_type == "litellm": - return LiteLLMEmbedding( - model_name=llm_settings.litellm_embedding_model, - api_key=llm_settings.litellm_embedding_api_key, - api_base=llm_settings.litellm_embedding_api_base, - ) - - raise Exception("embedding type is not supported !") -======= from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding from hugegraph_llm.models.embeddings.openai import OpenAIEmbedding from hugegraph_llm.models.embeddings.qianfan import QianFanEmbedding ->>>>>>> 38dce0b (feat(llm): vector db finished) class Embeddings: @@ -75,12 +39,6 @@ def get_embedding(self): if self.embedding_type == "ollama/local": assert llm_settings.ollama_embedding_model_dim, "ollama_embedding_model_dim is need" return OllamaEmbedding( -<<<<<<< HEAD - model_name=llm_settings.ollama_embedding_model, - host=llm_settings.ollama_embedding_host, - port=llm_settings.ollama_embedding_port, - ) -======= embedding_dimension=llm_settings.ollama_embedding_model_dim, model=llm_settings.ollama_embedding_model, host=llm_settings.ollama_embedding_host, @@ -93,17 +51,12 @@ def get_embedding(self): api_key=llm_settings.qianfan_embedding_api_key, secret_key=llm_settings.qianfan_embedding_secret_key, ) # type: ignore ->>>>>>> 38dce0b (feat(llm): vector db finished) if self.embedding_type == "litellm": return LiteLLMEmbedding( embedding_dimension=llm_settings.litellm_embedding_model_dim, model_name=llm_settings.litellm_embedding_model, api_key=llm_settings.litellm_embedding_api_key, api_base=llm_settings.litellm_embedding_api_base, -<<<<<<< HEAD - ) -======= ) # type: ignore ->>>>>>> 38dce0b (feat(llm): vector db finished) raise Exception("embedding type is not supported !") diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py index ac0522b99..27a0682aa 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py @@ -17,12 +17,7 @@ from typing import List, Optional -<<<<<<< HEAD -from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.log import log -======= from litellm import APIConnectionError, APIError, RateLimitError, aembedding, embedding ->>>>>>> 38dce0b (feat(llm): vector db finished) from tenacity import ( retry, retry_if_exception_type, @@ -45,9 +40,6 @@ def __init__( ) -> None: self.api_key = api_key self.api_base = api_base -<<<<<<< HEAD - self.model_name = model_name -======= self.model = model_name self.embedding_dimension = embedding_dimension @@ -55,7 +47,6 @@ def get_embedding_dim( self, ) -> int: return self.embedding_dimension ->>>>>>> 38dce0b (feat(llm): vector db finished) @retry( stop=stop_after_attempt(3), diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py index a29171ec5..02e0aa717 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py @@ -23,10 +23,6 @@ class OllamaEmbedding(BaseEmbedding): -<<<<<<< HEAD - def __init__(self, model_name: str, host: str = "127.0.0.1", port: int = 11434, **kwargs): - self.model_name = model_name -======= def __init__( self, model: str = "quentinz/bge-large-zh-v1.5", @@ -36,16 +32,10 @@ def __init__( **kwargs, ): self.model = model ->>>>>>> 38dce0b (feat(llm): vector db finished) self.client = ollama.Client(host=f"http://{host}:{port}", **kwargs) self.async_client = ollama.AsyncClient(host=f"http://{host}:{port}", **kwargs) self.embedding_dimension = embedding_dimension -<<<<<<< HEAD - def get_text_embedding(self, text: str) -> List[float]: - """Get embedding for a single text.""" - return self.get_texts_embeddings([text])[0] -======= def get_embedding_dim( self, ) -> int: @@ -54,13 +44,10 @@ def get_embedding_dim( def get_text_embedding(self, text: str) -> List[float]: """Comment""" return list(self.client.embed(model=self.model, input=text)["embeddings"][0]) ->>>>>>> 38dce0b (feat(llm): vector db finished) def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch. -<<<<<<< HEAD -======= This method efficiently processes multiple texts at once by leveraging Ollama's batching capabilities, which is more efficient than processing texts individually. @@ -70,7 +57,6 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: texts : List[str] A list of text strings to be embedded. ->>>>>>> 38dce0b (feat(llm): vector db finished) Returns ------- List[List[float]] @@ -87,27 +73,7 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: response = self.client.embed(model=self.model_name, input=texts)["embeddings"] return [list(inner_sequence) for inner_sequence in response] -<<<<<<< HEAD - async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts in a single batch asynchronously. - - Returns - ------- - List[List[float]] - A list of embedding vectors, where each vector is a list of floats. - The order of embeddings matches the order of input texts. - """ - if not hasattr(self.client, "embed"): - error_message = ( - "The required 'embed' method was not found on the Ollama client. " - "Please ensure your ollama library is up-to-date and supports batch embedding. " - ) - raise AttributeError(error_message) - response = await self.async_client.embed(model=self.model_name, input=texts) - return [list(inner_sequence) for inner_sequence in response["embeddings"]] -======= async def async_get_text_embedding(self, text: str) -> List[float]: """Comment""" response = await self.async_client.embeddings(model=self.model, prompt=text) return list(response["embedding"]) ->>>>>>> 38dce0b (feat(llm): vector db finished) diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index 6e30cca71..4c3793e3a 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -24,10 +24,7 @@ class OpenAIEmbedding: def __init__( self, -<<<<<<< HEAD -======= embedding_dimension: int = 1536, ->>>>>>> 38dce0b (feat(llm): vector db finished) model_name: str = "text-embedding-3-small", api_key: Optional[str] = None, api_base: Optional[str] = None, @@ -35,9 +32,6 @@ def __init__( api_key = api_key or "" self.client = OpenAI(api_key=api_key, base_url=api_base) self.aclient = AsyncOpenAI(api_key=api_key, base_url=api_base) -<<<<<<< HEAD - self.model_name = model_name -======= self.embedding_model_name = model_name self.embedding_dimension = embedding_dimension @@ -45,7 +39,6 @@ def get_embedding_dim( self, ) -> int: return self.embedding_dimension ->>>>>>> 38dce0b (feat(llm): vector db finished) def get_text_embedding(self, text: str) -> List[float]: """Comment""" diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py index 199384b12..aec5bc4df 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py @@ -118,15 +118,7 @@ async def agenerate_streaming( messages = [{"role": "user", "content": prompt}] try: -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - async_generator = await self.async_client.chat( - model=self.model, messages=messages, stream=True - ) -======= async_generator = await self.async_client.chat(model=self.model, messages=messages, stream=True) ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ======= async_generator = await self.async_client.chat(model=self.model, messages=messages, stream=True) >>>>>>> 8e0bf08 (chore: mark vectordb optional) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index 6eb805271..f2a7b0152 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -250,16 +250,7 @@ def run(self, **kwargs) -> Dict[str, Any]: :return: Final context after all operators have been executed. """ if len(self._operators) == 0: -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - self.extract_keywords().query_graphdb( - max_graph_items=kwargs.get("max_graph_items") - ).synthesize_answer() -======= self.extract_keywords().query_graphdb(max_graph_items=kwargs.get('max_graph_items')).synthesize_answer() ->>>>>>> 38dce0b (feat(llm): vector db finished) ======= self.extract_keywords().query_graphdb(max_graph_items=kwargs.get("max_graph_items")).synthesize_answer() >>>>>>> 87ee5d3 (style: format code with black line-length 120) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index 4f288a9de..72bb8a810 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -16,47 +16,19 @@ # under the License. -<<<<<<< HEAD -import asyncio -import os -from typing import Dict, Any, List - -from hugegraph_llm.config import resource_path, llm_settings, huge_settings -from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex -from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import ( - get_embeddings_parallel, - get_filename_prefix, - get_index_folder_name, -) -======= from typing import Any, Dict, List from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding ->>>>>>> 38dce0b (feat(llm): vector db finished) # FIXME: we need keep the logic same with build_semantic_index.py class BuildGremlinExampleIndex: -<<<<<<< HEAD - def __init__(self, embedding: BaseEmbedding, examples: List[Dict[str, str]]): - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "gremlin_examples")) - self.examples = examples - self.embedding = embedding - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(embedding, "model_name", None) - ) -======= def __init__(self, embedding: BaseEmbedding, examples: List[Dict[str, str]], vector_index: type[VectorStoreBase]): self.vector_index_name = "gremlin_examples" self.examples = examples self.embedding = embedding self.vector_index = vector_index ->>>>>>> 38dce0b (feat(llm): vector db finished) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: # !: We have assumed that self.example is not empty @@ -67,10 +39,6 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: if len(self.examples) > 0: vector_index = self.vector_index.from_name(embed_dim, self.vector_index_name) vector_index.add(examples_embedding, self.examples) -<<<<<<< HEAD - vector_index.to_index_file(self.index_dir, self.filename_prefix) -======= vector_index.save_index_by_name(self.vector_index_name) ->>>>>>> 38dce0b (feat(llm): vector db finished) context["embed_dim"] = embed_dim return context diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index b055cf62e..94ff2cee8 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -42,12 +42,6 @@ def __init__( self.num_examples = num_examples if not vector_index.exist("gremlin_examples"): log.warning("No gremlin example index found, will generate one.") -<<<<<<< HEAD - self.vector_index = vector_index.from_name(self.embedding.get_embedding_dim(), "gremlin_examples") -<<<<<<< HEAD - -======= ->>>>>>> 8e0bf08 (chore: mark vectordb optional) ======= self.vector_index = vector_index.from_name( self.embedding.get_embedding_dim(), "gremlin_examples" diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py index 97d2e8262..0553fad89 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py @@ -18,14 +18,8 @@ from typing import Any, Dict -<<<<<<< HEAD -<<<<<<< HEAD -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index import VectorIndex -======= from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) ======= from hugegraph_llm.config import huge_settings from hugegraph_llm.indices.vector_index.base import VectorStoreBase @@ -39,21 +33,7 @@ class VectorIndexQuery: def __init__(self, vector_index: type[VectorStoreBase], embedding: BaseEmbedding, topk: int = 3): self.embedding = embedding self.topk = topk -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(embedding, "model_name", None) - ) - self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) -======= self.vector_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "chunks") ->>>>>>> 38dce0b (feat(llm): vector db finished) ======= self.vector_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "chunks") >>>>>>> 8e0bf08 (chore: mark vectordb optional) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py index 13b23fa2c..ba9b93243 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py @@ -51,16 +51,7 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: llm_output = self.llm.generate(prompt=prompt) data["triples"] = [] extract_triples_by_regex(llm_output, data) -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - print( - f"LLM {self.__class__.__name__} input:{prompt} \n" - f" output: {llm_output} \n data: {data}" - ) -======= print(f"LLM {self.__class__.__name__} input:{prompt} \n" f" output: {llm_output} \n data: {data}") ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ======= print(f"LLM {self.__class__.__name__} input:{prompt} \n output: {llm_output} \n data: {data}") >>>>>>> 8e0bf08 (chore: mark vectordb optional) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index fd4583263..4ce9c6513 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -53,16 +53,7 @@ def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional return None example_strings = [] for example in examples: -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD - example_strings.append( - f"- query: {example['query']}\n" - f"- gremlin:\n```gremlin\n{example['gremlin']}\n```" - ) -======= example_strings.append(f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```") ->>>>>>> 38dce0b (feat(llm): vector db finished) ======= example_strings.append(f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```") >>>>>>> 8e0bf08 (chore: mark vectordb optional) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index 0fa1c0f04..7d0cbcc63 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -87,20 +87,12 @@ def _extract_keywords_from_response( for match in matches: match = match[len(start_token) :].strip() -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -======= ->>>>>>> 3aeef7d (fix) keywords.extend( k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1 ) -<<<<<<< HEAD -======= keywords.extend(k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1) ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ======= keywords.extend(k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1) >>>>>>> 8e0bf08 (chore: mark vectordb optional) 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 d517db8b9..fae3d7626 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 @@ -125,13 +125,7 @@ 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( -<<<<<<< HEAD -<<<<<<< HEAD - "Invalid property graph! No JSON object found, " - "please check the output format example in prompt." -======= "Invalid property graph! No JSON object found, " "please check the output format example in prompt." ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ======= "Invalid property graph! No JSON object found, please check the output format example in prompt." >>>>>>> 8e0bf08 (chore: mark vectordb optional) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py index 1e33514ca..972c828ec 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py @@ -86,14 +86,7 @@ def _extract_schema(self, response: str) -> Dict[str, Any]: raise RuntimeError("Invalid JSON response from LLM") from e def build_prompt( -<<<<<<< HEAD - self, - raw_texts: List[str], - query_examples: List[Dict[str, str]], - few_shot_schema: Dict[str, Any], -======= self, raw_texts: List[str], query_examples: List[Dict[str, str]], few_shot_schema: Dict[str, Any] ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ) -> str: return self.schema_prompt.format( raw_texts=self._format_raw_texts(raw_texts), diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index 6da7a6567..183341115 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -140,19 +140,11 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): elif filename == "vertices.json": data_full = client.gremlin().exec(query)["data"][0]["vertices"] data = ( -<<<<<<< HEAD -<<<<<<< HEAD -<<<<<<< HEAD -======= ->>>>>>> 3aeef7d (fix) [ {key: value for key, value in vertex.items() if key != "id"} for vertex in data_full ] -<<<<<<< HEAD -======= [{key: value for key, value in vertex.items() if key != "id"} for vertex in data_full] ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ======= [{key: value for key, value in vertex.items() if key != "id"} for vertex in data_full] >>>>>>> 8e0bf08 (chore: mark vectordb optional) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/log.py b/hugegraph-llm/src/hugegraph_llm/utils/log.py old mode 100755 new mode 100644 index aef36e449..3d956a51b --- a/hugegraph-llm/src/hugegraph_llm/utils/log.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/log.py @@ -27,15 +27,7 @@ # Initialize the root logger first with Rich handler root_logger = init_logger( -<<<<<<< HEAD - log_output=LOG_FILE, - log_level=INFO, - logger_name="root", - propagate_logs=True, - stdout_logging=True, -======= log_output=LOG_FILE, log_level=INFO, logger_name="root", propagate_logs=True, stdout_logging=True ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ) # Initialize custom logger diff --git a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py index 76fd4ccd1..734d87263 100644 --- a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py +++ b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py @@ -32,10 +32,4 @@ def test_stream_generate(self): def on_token_callback(chunk): print(chunk, end="", flush=True) -<<<<<<< HEAD - ollama_client.generate_streaming( - prompt="What is the capital of France?", on_token_callback=on_token_callback - ) -======= ollama_client.generate_streaming(prompt="What is the capital of France?", on_token_callback=on_token_callback) ->>>>>>> 87ee5d3 (style: format code with black line-length 120) diff --git a/hugegraph-python-client/src/pyhugegraph/api/graph.py b/hugegraph-python-client/src/pyhugegraph/api/graph.py index 94593836a..bce2acab4 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/graph.py +++ b/hugegraph-python-client/src/pyhugegraph/api/graph.py @@ -138,25 +138,11 @@ def addEdges(self, input_data) -> Optional[List[EdgeData]]: return None @router.http("PUT", "graph/edges/{edge_id}?action=append") -<<<<<<< HEAD - def appendEdge( - self, edge_id, properties # pylint: disable=unused-argument - ) -> Optional[EdgeData]: -======= - def appendEdge(self, edge_id, properties) -> Optional[EdgeData]: # pylint: disable=unused-argument ->>>>>>> 87ee5d3 (style: format code with black line-length 120) if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None @router.http("PUT", "graph/edges/{edge_id}?action=eliminate") -<<<<<<< HEAD - def eliminateEdge( - self, edge_id, properties # pylint: disable=unused-argument - ) -> Optional[EdgeData]: -======= - def eliminateEdge(self, edge_id, properties) -> Optional[EdgeData]: # pylint: disable=unused-argument ->>>>>>> 87ee5d3 (style: format code with black line-length 120) if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None diff --git a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py index 5026dce21..de2bb67ee 100644 --- a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py +++ b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py @@ -25,17 +25,9 @@ schema.propertyKey("name").asText().ifNotExist().create() schema.propertyKey("birthDate").asText().ifNotExist().create() schema.vertexLabel("Person").properties("name", "birthDate").usePrimaryKeyId().primaryKeys( -<<<<<<< HEAD - "name" - ).ifNotExist().create() - schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys( - "name" - ).ifNotExist().create() -======= "name" ).ifNotExist().create() schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys("name").ifNotExist().create() ->>>>>>> 87ee5d3 (style: format code with black line-length 120) schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() print(schema.getVertexLabels()) diff --git a/hugegraph-python-client/src/pyhugegraph/utils/util.py b/hugegraph-python-client/src/pyhugegraph/utils/util.py index 6df64deff..3fe300092 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/util.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/util.py @@ -58,14 +58,6 @@ def check_if_success(response, error=None): req_body = req.body if req.body else "Empty body" response_body = response.text if response.text else "Empty body" log.error( -<<<<<<< HEAD - "Error-Client: Request URL: %s, Request Body: %s, Response Body: %s", - req.url, - req_body, - response_body, -======= - "Error-Client: Request URL: %s, Request Body: %s, Response Body: %s", req.url, req_body, response_body ->>>>>>> 87ee5d3 (style: format code with black line-length 120) ) raise error return True From 0a2e53d89c2f0d16b922983da2cd32c6705292e4 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Thu, 9 Oct 2025 15:36:12 +0800 Subject: [PATCH 32/71] fix: remove remaining conflict markers in operators and utils --- .../src/hugegraph_llm/demo/rag_demo/admin_block.py | 8 -------- .../src/hugegraph_llm/document/chunk_split.py | 10 ---------- hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py | 6 ------ .../src/hugegraph_llm/operators/graph_rag_task.py | 9 --------- 4 files changed, 33 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 00da2a973..7eecefb9e 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -109,14 +109,6 @@ def create_admin_block(): # Error message box, initially hidden error_message = gr.Textbox(label="", visible=False, interactive=False, elem_classes="error-message") -======= - error_message = gr.Textbox(label="", visible=False, interactive=False, elem_classes="error-message") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= - error_message = gr.Textbox( - label="", visible=False, interactive=False, elem_classes="error-message" - ) ->>>>>>> 3aeef7d (fix) # Button to submit password submit_button = gr.Button("Submit") diff --git a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py index c9ece77bb..4c4a42e3a 100644 --- a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py +++ b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py @@ -31,15 +31,6 @@ def __init__( else: raise ValueError("Argument `language` must be zh or en!") if split_type == "paragraph": - self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=30, separators=separators) - elif split_type == "sentence": - self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=50, chunk_overlap=0, separators=separators) -======= - self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=30, separators=separators) - elif split_type == "sentence": - self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=50, chunk_overlap=0, separators=separators) ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=500, chunk_overlap=30, separators=separators ) @@ -47,7 +38,6 @@ def __init__( self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=50, chunk_overlap=0, separators=separators ) ->>>>>>> 3aeef7d (fix) else: raise ValueError("Arg `type` must be paragraph, sentence!") diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py index aec5bc4df..6d08ce8cd 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py @@ -118,15 +118,9 @@ async def agenerate_streaming( messages = [{"role": "user", "content": prompt}] try: - async_generator = await self.async_client.chat(model=self.model, messages=messages, stream=True) -======= - async_generator = await self.async_client.chat(model=self.model, messages=messages, stream=True) ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= async_generator = await self.async_client.chat( model=self.model, messages=messages, stream=True ) ->>>>>>> 3aeef7d (fix) async for chunk in async_generator: token = chunk.get("message", {}).get("content", "") if on_token_callback: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index f2a7b0152..ceadecbd4 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -250,18 +250,9 @@ def run(self, **kwargs) -> Dict[str, Any]: :return: Final context after all operators have been executed. """ if len(self._operators) == 0: - self.extract_keywords().query_graphdb(max_graph_items=kwargs.get('max_graph_items')).synthesize_answer() -======= - self.extract_keywords().query_graphdb(max_graph_items=kwargs.get("max_graph_items")).synthesize_answer() ->>>>>>> 87ee5d3 (style: format code with black line-length 120) -======= - self.extract_keywords().query_graphdb(max_graph_items=kwargs.get("max_graph_items")).synthesize_answer() ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= self.extract_keywords().query_graphdb( max_graph_items=kwargs.get("max_graph_items") ).synthesize_answer() ->>>>>>> 3aeef7d (fix) context = kwargs From 3c6a9eef76bf7c49ab03068ad2f7572ecc0dafb0 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Thu, 9 Oct 2025 15:45:54 +0800 Subject: [PATCH 33/71] fix: finalize conflict marker cleanup in llm operators and utils --- .../index_op/gremlin_example_index_query.py | 2 -- .../operators/index_op/vector_index_query.py | 25 ------------------- 2 files changed, 27 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index 94ff2cee8..3f660880d 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -42,11 +42,9 @@ def __init__( self.num_examples = num_examples if not vector_index.exist("gremlin_examples"): log.warning("No gremlin example index found, will generate one.") -======= self.vector_index = vector_index.from_name( self.embedding.get_embedding_dim(), "gremlin_examples" ) ->>>>>>> 3aeef7d (fix) self._build_default_example_index() else: self.vector_index = vector_index.from_name( diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py index 0553fad89..0ae1dd4f6 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py @@ -18,14 +18,9 @@ from typing import Any, Dict -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex -======= from hugegraph_llm.config import huge_settings from hugegraph_llm.indices.vector_index.base import VectorStoreBase ->>>>>>> 38dce0b (feat(llm): vector db finished) from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import get_filename_prefix, get_index_folder_name from hugegraph_llm.utils.log import log @@ -33,33 +28,13 @@ class VectorIndexQuery: def __init__(self, vector_index: type[VectorStoreBase], embedding: BaseEmbedding, topk: int = 3): self.embedding = embedding self.topk = topk - self.vector_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "chunks") -======= - self.vector_index = vector_index.from_name(embedding.get_embedding_dim(), huge_settings.graph_name, "chunks") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= self.vector_index = vector_index.from_name( embedding.get_embedding_dim(), huge_settings.graph_name, "chunks" ) ->>>>>>> 3aeef7d (fix) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query = context.get("query") query_embedding = self.embedding.get_texts_embeddings([query])[0] -======= - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(embedding, "model_name", None) - ) - self.vector_index = FaissVectorIndex.from_index_file(self.index_dir, self.filename_prefix) - - def run(self, context: Dict[str, Any]) -> Dict[str, Any]: - query = context.get("query") - query_embedding = self.embedding.get_texts_embeddings([query])[0] ->>>>>>> 902fee5 (feat(llm): some type bug && revert to FaissVectorIndex) # TODO: why set dis_threshold=2? results = self.vector_index.search(query_embedding, self.topk, dis_threshold=2) # TODO: check format results From 895ba476675b35891582d940fe3dabf88d3a0838 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Thu, 9 Oct 2025 18:03:42 +0800 Subject: [PATCH 34/71] fix --- .../hugegraph_llm/operators/llm_op/disambiguate_data.py | 8 -------- .../hugegraph_llm/operators/llm_op/gremlin_generate.py | 6 ------ .../hugegraph_llm/operators/llm_op/keyword_extract.py | 6 ------ .../operators/llm_op/property_graph_extract.py | 3 --- hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py | 9 --------- hugegraph-python-client/src/pyhugegraph/api/graph.py | 2 ++ 6 files changed, 2 insertions(+), 32 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py index ba9b93243..5913ea307 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py @@ -51,15 +51,7 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: llm_output = self.llm.generate(prompt=prompt) data["triples"] = [] extract_triples_by_regex(llm_output, data) - print(f"LLM {self.__class__.__name__} input:{prompt} \n" f" output: {llm_output} \n data: {data}") -======= print(f"LLM {self.__class__.__name__} input:{prompt} \n output: {llm_output} \n data: {data}") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= - print( - f"LLM {self.__class__.__name__} input:{prompt} \n output: {llm_output} \n data: {data}" - ) ->>>>>>> 3aeef7d (fix) data["call_count"] = data.get("call_count", 0) + 1 return data diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index 4ce9c6513..39f49dbdd 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -53,15 +53,9 @@ def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional return None example_strings = [] for example in examples: - example_strings.append(f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```") -======= - example_strings.append(f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```") ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= example_strings.append( f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```" ) ->>>>>>> 3aeef7d (fix) return "\n\n".join(example_strings) def _format_vertices(self, vertices: Optional[List[str]]) -> Optional[str]: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index 7d0cbcc63..425f2a70b 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -92,12 +92,6 @@ def _extract_keywords_from_response( for k in re.split(r"[,，]+", match) if len(k.strip()) > 1 ) - keywords.extend(k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1) -======= - keywords.extend(k.lower() if lowercase else k for k in re.split(r"[,，]+", match) if len(k.strip()) > 1) ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= ->>>>>>> 3aeef7d (fix) # if the keyword consists of multiple words, split into sub-words (removing stopwords) results = set(keywords) 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 fae3d7626..f4933d676 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 @@ -125,10 +125,7 @@ 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 object found, " "please check the output format example in prompt." -======= "Invalid property graph! No JSON object found, please check the output format example in prompt." ->>>>>>> 8e0bf08 (chore: mark vectordb optional) ) return [] json_str = json_match.group(1).strip() diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index 183341115..6965b6239 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -140,16 +140,7 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): elif filename == "vertices.json": data_full = client.gremlin().exec(query)["data"][0]["vertices"] data = ( - [ - {key: value for key, value in vertex.items() if key != "id"} - for vertex in data_full - ] [{key: value for key, value in vertex.items() if key != "id"} for vertex in data_full] -======= - [{key: value for key, value in vertex.items() if key != "id"} for vertex in data_full] ->>>>>>> 8e0bf08 (chore: mark vectordb optional) -======= ->>>>>>> 3aeef7d (fix) if all_pk_flag else data_full ) diff --git a/hugegraph-python-client/src/pyhugegraph/api/graph.py b/hugegraph-python-client/src/pyhugegraph/api/graph.py index bce2acab4..2372b6522 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/graph.py +++ b/hugegraph-python-client/src/pyhugegraph/api/graph.py @@ -138,11 +138,13 @@ def addEdges(self, input_data) -> Optional[List[EdgeData]]: return None @router.http("PUT", "graph/edges/{edge_id}?action=append") + def appendEdge(self, edge_id, properties): # pylint: disable=unused-argument if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None @router.http("PUT", "graph/edges/{edge_id}?action=eliminate") + def eliminateEdge(self, edge_id, properties): # pylint: disable=unused-argument if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None From e3d4ee1c365fe8b22a91f0f269c75274248db630 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 10 Oct 2025 19:39:16 +0800 Subject: [PATCH 35/71] feat(embeddings): drop QianFan provider and remove related configs - Remove qianfan.py and references from init_embedding - Clean LLMConfig: delete QianFan/WenXin fields - No behavior change for other providers --- .../src/hugegraph_llm/config/llm_config.py | 24 +------ .../models/embeddings/init_embedding.py | 8 --- .../models/embeddings/qianfan.py | 66 ------------------- 3 files changed, 1 insertion(+), 97 deletions(-) delete mode 100644 hugegraph-llm/src/hugegraph_llm/models/embeddings/qianfan.py diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index af22b71b0..75494ab31 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -70,30 +70,8 @@ class LLMConfig(BaseConfig): _env_ollama_dim = os.getenv("OLLAMA_EMBEDDING_MODEL_DIM") ollama_embedding_model_dim: Optional[int] = int(_env_ollama_dim) if _env_ollama_dim else None - # 4. QianFan/WenXin settings - # TODO: update to one token key mode - qianfan_chat_api_key: Optional[str] = None - qianfan_chat_secret_key: Optional[str] = None - qianfan_chat_access_token: Optional[str] = None - qianfan_extract_api_key: Optional[str] = None - qianfan_extract_secret_key: Optional[str] = None - qianfan_extract_access_token: Optional[str] = None - qianfan_text2gql_api_key: Optional[str] = None - qianfan_text2gql_secret_key: Optional[str] = None - qianfan_text2gql_access_token: Optional[str] = None - qianfan_embedding_api_key: Optional[str] = None - qianfan_embedding_secret_key: Optional[str] = None - # 4.1 URL settings - qianfan_url_prefix: str = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop" - qianfan_chat_url: str = qianfan_url_prefix + "/chat/" - qianfan_chat_language_model: str = "ERNIE-Speed-128K" - qianfan_extract_language_model: str = "ERNIE-Speed-128K" - qianfan_text2gql_language_model: str = "ERNIE-Speed-128K" - qianfan_embed_url: str = qianfan_url_prefix + "/embeddings/" - qianfan_embedding_model_dim: int = 384 + # 4. QianFan/WenXin settings (removed) - # refer https://cloud.baidu.com/doc/WENXINWORKSHOP/s/alj562vvu to get more details - qianfan_embedding_model: str = "embedding-v1" # 5. LiteLLM settings litellm_chat_api_key: Optional[str] = None litellm_chat_api_base: Optional[str] = None diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index 143d54229..dfd23ccfe 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -20,7 +20,6 @@ from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding from hugegraph_llm.models.embeddings.openai import OpenAIEmbedding -from hugegraph_llm.models.embeddings.qianfan import QianFanEmbedding class Embeddings: @@ -44,13 +43,6 @@ def get_embedding(self): host=llm_settings.ollama_embedding_host, port=llm_settings.ollama_embedding_port, ) - if self.embedding_type == "qianfan_wenxin": - return QianFanEmbedding( - embedding_dimension=llm_settings.litellm_embedding_model_dim, - model_name=llm_settings.qianfan_embedding_model, - api_key=llm_settings.qianfan_embedding_api_key, - secret_key=llm_settings.qianfan_embedding_secret_key, - ) # type: ignore if self.embedding_type == "litellm": return LiteLLMEmbedding( embedding_dimension=llm_settings.litellm_embedding_model_dim, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/qianfan.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/qianfan.py deleted file mode 100644 index e5d5463ef..000000000 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/qianfan.py +++ /dev/null @@ -1,66 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - - -from typing import Optional, List - -import qianfan - -from hugegraph_llm.config import llm_settings - -""" -"QianFan" platform can be understood as a unified LLM platform that encompasses the -WenXin large model along with other -common open-source models. - -It enables the invocation and switching between WenXin and these open-source models. -""" - - -class QianFanEmbedding: - def __init__( - self, - embedding_dimension: int, - model_name: str = "embedding-v1", - api_key: Optional[str] = None, - secret_key: Optional[str] = None, - ): - qianfan.get_config().AK = api_key or llm_settings.qianfan_embedding_api_key - qianfan.get_config().SK = secret_key or llm_settings.qianfan_embedding_secret_key - self.embedding_model_name = model_name - self.client = qianfan.Embedding() - self.embedding_dimension = embedding_dimension - - def get_embedding_dim( - self, - ) -> int: - return self.embedding_dimension - - def get_text_embedding(self, text: str) -> List[float]: - """Usage refer: https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlmokk9qn""" - response = self.client.do(model=self.embedding_model_name, texts=[text]) - return response["body"]["data"][0]["embedding"] - - def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Usage refer: https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlmokk9qn""" - response = self.client.do(model=self.embedding_model_name, texts=texts) - return [data["embedding"] for data in response["body"]["data"]] - - async def async_get_text_embedding(self, text: str) -> List[float]: - """Usage refer: https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlmokk9qn""" - response = await self.client.ado(model=self.embedding_model_name, texts=[text]) - return response["body"]["data"][0]["embedding"] From a4cfe05e4f4d4f07e2aeba0b619c6f465edcbfde Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 10 Oct 2025 19:39:28 +0800 Subject: [PATCH 36/71] feat(vector-index): UI config for Milvus/Qdrant and CUR_VECTOR_INDEX env - Add backend connection test & persistence in Gradio UI - Support CUR_VECTOR_INDEX env override in IndexConfig - Lazy import Milvus/Qdrant with user-friendly errors --- .../src/hugegraph_llm/config/index_config.py | 2 +- .../demo/rag_demo/configs_block.py | 115 +++++++++++++++--- .../hugegraph_llm/utils/vector_index_utils.py | 40 ++++-- 3 files changed, 130 insertions(+), 27 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/config/index_config.py b/hugegraph-llm/src/hugegraph_llm/config/index_config.py index ad0db5975..63895e6a7 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/index_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/index_config.py @@ -35,4 +35,4 @@ class IndexConfig(BaseConfig): milvus_user: str = os.environ.get("MILVUS_USER", "") milvus_password: str = os.environ.get("MILVUS_PASSWORD", "") - cur_vector_index: str = "Faiss" + cur_vector_index: str = os.environ.get("CUR_VECTOR_INDEX", "Faiss") diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 112892e7c..6978158c8 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -104,24 +104,76 @@ def test_api_connection( except (json.decoder.JSONDecodeError, AttributeError) as e: raise gr.Error(resp.text) from e return resp.status_code +def apply_vector_engine(engine: str): + # Persist the vector engine selection + setattr(index_settings, "cur_vector_index", engine) + try: + index_settings.update_env() + except Exception: # pylint: disable=W0718 + pass + gr.Info("Configured!") -def config_qianfan_model(arg1, arg2, arg3=None, settings_prefix=None, origin_call=None) -> int: - setattr(llm_settings, f"qianfan_{settings_prefix}_api_key", arg1) - setattr(llm_settings, f"qianfan_{settings_prefix}_secret_key", arg2) - if arg3: - setattr(llm_settings, f"qianfan_{settings_prefix}_language_model", arg3) - params = { - "grant_type": "client_credentials", - "client_id": arg1, - "client_secret": arg2, - } - status_code = test_api_connection( - "https://aip.baidubce.com/oauth/2.0/token", - "POST", - params=params, - origin_call=origin_call, - ) +def apply_vector_engine_backend( + engine: str, + host: Optional[str] = None, + port: Optional[str] = None, + user: Optional[str] = None, + password: Optional[str] = None, + api_key: Optional[str] = None, + origin_call=None, +) -> int: + """Test connection and persist per-engine connection settings""" + status_code = -1 + + # Test connection first + try: + if engine == "Milvus": + from pymilvus import connections, utility + connections.connect(host=host, port=int(port or 19530), user=user or "", password=password or "") + # Test if we can list collections + _ = utility.list_collections() + connections.disconnect("default") + status_code = 200 + elif engine == "Qdrant": + from qdrant_client import QdrantClient + client = QdrantClient(host=host, port=int(port or 6333), api_key=api_key) + # Test if we can get collections + _ = client.get_collections() + status_code = 200 + except ImportError as e: + msg = f"Missing dependency: {e}. Please install with: uv sync --extra vectordb" + if origin_call is None: + raise gr.Error(msg) from e + return -1 + except Exception as e: + msg = f"Connection failed: {e}" + log.error(msg) + if origin_call is None: + raise gr.Error(msg) from e + return -1 + + # Persist settings after successful test + if engine == "Milvus": + if host is not None: + index_settings.milvus_host = host + if port is not None and str(port).strip(): + index_settings.milvus_port = int(port) # type: ignore[arg-type] + index_settings.milvus_user = user or "" + index_settings.milvus_password = password or "" + elif engine == "Qdrant": + if host is not None: + index_settings.qdrant_host = host + if port is not None and str(port).strip(): + index_settings.qdrant_port = int(port) # type: ignore[arg-type] + # Empty string treated as None for api key + index_settings.qdrant_api_key = api_key or None + + try: + index_settings.update_env() + except Exception: # pylint: disable=W0718 + pass + gr.Info("Configured!") return status_code @@ -688,6 +740,37 @@ def reranker_settings(reranker_type): inputs=[engine_selector], ) + @gr.render(inputs=[engine_selector]) + def vector_engine_settings(engine): + if engine == "Milvus": + with gr.Row(): + milvus_inputs = [ + gr.Textbox(value=index_settings.milvus_host, label="host"), + gr.Textbox(value=str(index_settings.milvus_port), label="port"), + gr.Textbox(value=index_settings.milvus_user, label="user"), + gr.Textbox(value=index_settings.milvus_password, label="password", type="password"), + ] + apply_backend_button = gr.Button("Apply Configuration") + apply_backend_button.click( + partial(apply_vector_engine_backend, "Milvus"), inputs=milvus_inputs + ) + elif engine == "Qdrant": + with gr.Row(): + qdrant_inputs = [ + gr.Textbox(value=index_settings.qdrant_host, label="host"), + gr.Textbox(value=str(index_settings.qdrant_port), label="port"), + gr.Textbox(value=(index_settings.qdrant_api_key or ""), label="api_key", type="password"), + ] + apply_backend_button = gr.Button("Apply Configuration") + apply_backend_button.click( + lambda h, p, k: apply_vector_engine_backend("Qdrant", h, p, None, None, k), + inputs=qdrant_inputs, + ) + else: + gr.Markdown("✅ Faiss 本地索引无需额外配置。") + apply_faiss_button = gr.Button("Apply Configuration") + apply_faiss_button.click(lambda: apply_vector_engine(engine)) + # The reason for returning this partial value is the functional need to refresh the ui return graph_config_input diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 4f56db636..9468f823b 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -24,8 +24,6 @@ from hugegraph_llm.config import huge_settings, index_settings from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex -from hugegraph_llm.indices.vector_index.milvus_vector_store import MilvusVectorIndex -from hugegraph_llm.indices.vector_index.qdrant_vector_store import QdrantVectorIndex from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.models.llms.init_llm import LLMs from hugegraph_llm.operators.kg_construction_task import KgBuilder @@ -93,11 +91,33 @@ def build_vector_index(input_file, input_text): def get_vector_index_class(vector_index_str: str) -> Type[VectorStoreBase]: - mapping = { - "Faiss": FaissVectorIndex, - "Milvus": MilvusVectorIndex, - "Qdrant": QdrantVectorIndex, - } - ret = mapping.get(vector_index_str) - assert ret - return ret # type: ignore + if vector_index_str == "Faiss": + return FaissVectorIndex # type: ignore[return-value] + if vector_index_str == "Milvus": + try: + from hugegraph_llm.indices.vector_index.milvus_vector_store import ( # pylint: disable=import-outside-toplevel + MilvusVectorIndex, + ) + + return MilvusVectorIndex # type: ignore[return-value] + except Exception as e: # pylint: disable=broad-except + raise gr.Error( + f"Milvus engine selected but dependency not available: {e}.\n" + "Fix it by running: 'uv sync --extra vectordb' (recommended) or install 'pymilvus' manually.\n" + "Alternatively, switch vector engine to Faiss/Qdrant in the UI." + ) + if vector_index_str == "Qdrant": + try: + from hugegraph_llm.indices.vector_index.qdrant_vector_store import ( # pylint: disable=import-outside-toplevel + QdrantVectorIndex, + ) + + return QdrantVectorIndex # type: ignore[return-value] + except Exception as e: # pylint: disable=broad-except + raise gr.Error( + f"Qdrant engine selected but dependency not available: {e}.\n" + "Fix it by running: 'uv sync --extra vectordb' (recommended) or install 'qdrant-client' manually.\n" + "Alternatively, switch vector engine to Faiss/Milvus in the UI." + ) + # Fallback to Faiss + return FaissVectorIndex # type: ignore[return-value] From 6e31b8739f067ae076c38ed0416296b39b8865ba Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 10 Oct 2025 19:39:42 +0800 Subject: [PATCH 37/71] refactor(nodes): unify embedding initialization via Embeddings().get_embedding() --- .../src/hugegraph_llm/nodes/common_node/merge_rerank_node.py | 4 ++-- .../hugegraph_llm/nodes/index_node/build_semantic_index.py | 4 ++-- .../src/hugegraph_llm/nodes/index_node/build_vector_index.py | 4 ++-- .../nodes/index_node/gremlin_example_index_query.py | 4 ++-- .../hugegraph_llm/nodes/index_node/semantic_id_query_node.py | 4 ++-- .../src/hugegraph_llm/nodes/index_node/vector_query_node.py | 4 ++-- 6 files changed, 12 insertions(+), 12 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py index 78f53e231..405be1bf0 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py @@ -16,7 +16,7 @@ from typing import Dict, Any from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.common_op.merge_dedup_rerank import MergeDedupRerank -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.config import huge_settings, llm_settings from hugegraph_llm.utils.log import log @@ -34,7 +34,7 @@ def node_init(self): """ try: # Read user configuration parameters from wk_input - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() graph_ratio = self.wk_input.graph_ratio or 0.5 rerank_method = self.wk_input.rerank_method or "bleu" near_neighbor_first = self.wk_input.near_neighbor_first or False 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 c01cffc91..327e3cea8 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 @@ -14,7 +14,7 @@ # limitations under the License. from hugegraph_llm.config import llm_settings -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -26,7 +26,7 @@ class BuildSemanticIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): - self.build_semantic_index_op = BuildSemanticIndex(get_embedding(llm_settings)) + self.build_semantic_index_op = BuildSemanticIndex(Embeddings().get_embedding()) return super().node_init() def operator_schedule(self, data_json): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py index 1f6a3c75b..a58a0dda4 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py @@ -14,7 +14,7 @@ # limitations under the License. from hugegraph_llm.config import llm_settings -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -26,7 +26,7 @@ class BuildVectorIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): - self.build_vector_index_op = BuildVectorIndex(get_embedding(llm_settings)) + self.build_vector_index_op = BuildVectorIndex(Embeddings().get_embedding()) return super().node_init() def operator_schedule(self, data_json): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py index e9283598a..88b019e33 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -24,7 +24,7 @@ from hugegraph_llm.operators.index_op.gremlin_example_index_query import ( GremlinExampleIndexQuery, ) -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings class GremlinExampleIndexQueryNode(BaseNode): @@ -32,7 +32,7 @@ class GremlinExampleIndexQueryNode(BaseNode): def node_init(self): # Build operator (index lazy-loading handled in operator) - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() example_num = getattr(self.wk_input, "example_num", None) if not isinstance(example_num, int): example_num = 2 diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py index bf605aa49..db535bd6f 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -17,7 +17,7 @@ from typing import Dict, Any from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.semantic_id_query import SemanticIdQuery -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.config import huge_settings, llm_settings from hugegraph_llm.utils.log import log @@ -38,7 +38,7 @@ def node_init(self): if not graph_name: return CStatus(-1, "graph_name is required in wk_input") - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() by = self.wk_input.semantic_by or "keywords" topk_per_keyword = ( self.wk_input.topk_per_keyword or huge_settings.topk_per_keyword diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py index 48b50acf3..d4c5c8026 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py @@ -17,7 +17,7 @@ from hugegraph_llm.config import llm_settings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.utils.log import log @@ -34,7 +34,7 @@ def node_init(self): """ try: # 从 wk_input 中读取用户配置参数 - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() max_items = ( self.wk_input.max_items if self.wk_input.max_items is not None else 3 ) From b446e4fc0bb4d2ac0fe779f14920195106dc8622 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 10 Oct 2025 19:39:57 +0800 Subject: [PATCH 38/71] refactor(flow): use FaissVectorIndex + Embeddings for graph index info --- .../flows/get_graph_index_info.py | 25 ++++++++----------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py index 7d2735352..6547c67b3 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py @@ -16,10 +16,10 @@ import json import os -from hugegraph_llm.config import huge_settings, llm_settings, resource_path +from hugegraph_llm.config import huge_settings, resource_path from hugegraph_llm.flows.common import BaseFlow -from hugegraph_llm.indices.vector_index import VectorIndex -from hugegraph_llm.models.embeddings.init_embedding import model_map +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState from hugegraph_llm.nodes.hugegraph_node.fetch_graph_data import FetchGraphDataNode from PyCGraph import GPipeline @@ -27,6 +27,7 @@ get_filename_prefix, get_index_folder_name, ) +from hugegraph_llm.models.embeddings.init_embedding import Embeddings class GetGraphIndexInfoFlow(BaseFlow): @@ -49,18 +50,14 @@ def build_flow(self, *args, **kwargs): def post_deal(self, pipeline=None): graph_summary_info = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, - model_map.get(llm_settings.embedding_type, None), - ) - try: - vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) - except FileNotFoundError: + if not FaissVectorIndex.exist(folder_name, "graph_vids"): return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) + embed_dim = Embeddings().get_embedding().get_embedding_dim() + vector_index = FaissVectorIndex.from_name(embed_dim, folder_name, "graph_vids") + vector_index_info = vector_index.get_vector_index_info() graph_summary_info["vid_index"] = { - "embed_dim": vector_index.index.d, - "num_vectors": vector_index.index.ntotal, - "num_vids": len(vector_index.properties), + "embed_dim": vector_index_info["embed_dim"], + "num_vectors": vector_index_info["vector_info"]["chunk_vector_num"], + "num_vids": vector_index_info["vector_info"]["graph_properties_vector_num"], } return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) From 2df53b65e07d1ac43a26ef4208e6515417b127e9 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 10 Oct 2025 19:40:17 +0800 Subject: [PATCH 39/71] fix(embeddings): align wrappers with BaseEmbedding and async batch - LiteLLM/OpenAI/Ollama use unified 'model' field - Add async batch method usage in gremlin example index builder - Improve Ollama async batch fallback --- .../hugegraph_llm/models/embeddings/litellm.py | 16 ++++++---------- .../hugegraph_llm/models/embeddings/ollama.py | 13 ++++++++----- .../hugegraph_llm/models/embeddings/openai.py | 15 ++++++++++----- .../index_op/build_gremlin_example_index.py | 2 ++ 4 files changed, 26 insertions(+), 20 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py index 27a0682aa..3eb47c566 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py @@ -18,14 +18,10 @@ from typing import List, Optional from litellm import APIConnectionError, APIError, RateLimitError, aembedding, embedding -from tenacity import ( - retry, - retry_if_exception_type, - stop_after_attempt, - wait_exponential, -) +from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential -from litellm import embedding, RateLimitError, APIError, APIConnectionError, aembedding +from hugegraph_llm.models.embeddings.base import BaseEmbedding +from hugegraph_llm.utils.log import log class LiteLLMEmbedding(BaseEmbedding): @@ -57,7 +53,7 @@ def get_text_embedding(self, text: str) -> List[float]: """Get embedding for a single text.""" try: response = embedding( - model=self.model_name, + model=self.model, input=text, api_key=self.api_key, api_base=self.api_base, @@ -72,7 +68,7 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts.""" try: response = embedding( - model=self.model_name, + model=self.model, input=texts, api_key=self.api_key, api_base=self.api_base, @@ -87,7 +83,7 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] """Get embedding for a single text asynchronously.""" try: response = await aembedding( - model=self.model_name, + model=self.model, input=texts, api_key=self.api_key, api_base=self.api_base, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py index 02e0aa717..ac07836e1 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py @@ -70,10 +70,13 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: ) raise AttributeError(error_message) - response = self.client.embed(model=self.model_name, input=texts)["embeddings"] + response = self.client.embed(model=self.model, input=texts)["embeddings"] return [list(inner_sequence) for inner_sequence in response] - async def async_get_text_embedding(self, text: str) -> List[float]: - """Comment""" - response = await self.async_client.embeddings(model=self.model, prompt=text) - return list(response["embedding"]) + async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: + # Ollama python client may not provide batch async embeddings; fallback per item + results: List[List[float]] = [] + for t in texts: + response = await self.async_client.embeddings(model=self.model, prompt=t) + results.append(list(response["embedding"])) + return results diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index 4c3793e3a..135f71000 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -19,9 +19,10 @@ from typing import Optional, List from openai import OpenAI, AsyncOpenAI +from hugegraph_llm.models.embeddings.base import BaseEmbedding -class OpenAIEmbedding: +class OpenAIEmbedding(BaseEmbedding): def __init__( self, embedding_dimension: int = 1536, @@ -32,7 +33,7 @@ def __init__( api_key = api_key or "" self.client = OpenAI(api_key=api_key, base_url=api_base) self.aclient = AsyncOpenAI(api_key=api_key, base_url=api_base) - self.embedding_model_name = model_name + self.model = model_name self.embedding_dimension = embedding_dimension def get_embedding_dim( @@ -42,7 +43,7 @@ def get_embedding_dim( def get_text_embedding(self, text: str) -> List[float]: """Comment""" - response = self.client.embeddings.create(input=text, model=self.model_name) + response = self.client.embeddings.create(input=text, model=self.model) return response.data[0].embedding def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: @@ -63,7 +64,7 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: A list of embedding vectors, where each vector is a list of floats. The order of embeddings matches the order of input texts. """ - response = self.client.embeddings.create(input=texts, model=self.model_name) + response = self.client.embeddings.create(input=texts, model=self.model) return [data.embedding for data in response.data] async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: @@ -84,5 +85,9 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] A list of embedding vectors, where each vector is a list of floats. The order of embeddings should match the order of input texts. """ - response = await self.aclient.embeddings.create(input=texts, model=self.model_name) + response = await self.aclient.embeddings.create(input=texts, model=self.model) return [data.embedding for data in response.data] + + async def async_get_text_embedding(self, text: str) -> List[float]: + response = await self.aclient.embeddings.create(input=[text], model=self.model) + return response.data[0].embedding diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index 72bb8a810..d052f2f6e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -16,10 +16,12 @@ # under the License. +import asyncio from typing import Any, Dict, List from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding +from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel # FIXME: we need keep the logic same with build_semantic_index.py From 1893d6485549068a47c165266f99ba07e20c130d Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 10 Oct 2025 19:48:00 +0800 Subject: [PATCH 40/71] style: apply black formatting with line-length 100 - Reformat 42 files to comply with Black code style - Align with CI workflow requirements - No functional changes --- .../src/hugegraph_llm/api/admin_api.py | 4 ++- .../src/hugegraph_llm/config/generate.py | 8 ++++- .../demo/rag_demo/admin_block.py | 4 ++- .../demo/rag_demo/configs_block.py | 22 ++++++++++--- .../demo/rag_demo/other_block.py | 4 ++- .../src/hugegraph_llm/document/chunk_split.py | 4 ++- .../src/hugegraph_llm/flows/common.py | 4 +-- .../flows/rag_flow_graph_only.py | 16 +++------- .../flows/rag_flow_graph_vector.py | 28 ++++------------ .../src/hugegraph_llm/flows/rag_flow_raw.py | 4 +-- .../flows/rag_flow_vector_only.py | 8 ++--- .../src/hugegraph_llm/flows/scheduler.py | 4 +-- .../vector_index/qdrant_vector_store.py | 8 +++-- .../nodes/hugegraph_node/graph_query_node.py | 8 ++--- .../index_node/gremlin_example_index_query.py | 4 +-- .../index_node/semantic_id_query_node.py | 8 ++--- .../nodes/index_node/vector_query_node.py | 4 +-- .../nodes/llm_node/answer_synthesize_node.py | 4 +-- .../nodes/llm_node/keyword_extract_node.py | 10 ++---- .../nodes/llm_node/text2gremlin.py | 4 +-- .../index_op/build_gremlin_example_index.py | 7 +++- .../index_op/gremlin_example_index_query.py | 9 ++++-- .../operators/index_op/vector_index_query.py | 4 ++- .../operators/llm_op/disambiguate_data.py | 4 ++- .../operators/llm_op/schema_build.py | 5 ++- .../llm_op/unstructured_data_utils.py | 5 ++- .../src/hugegraph_llm/state/ai_state.py | 6 +--- .../hugegraph_llm/utils/hugegraph_utils.py | 5 ++- hugegraph-llm/src/hugegraph_llm/utils/log.py | 6 +++- .../tests/models/llms/test_ollama_client.py | 4 ++- .../src/pyhugegraph/api/auth.py | 20 +++++++++--- .../src/pyhugegraph/api/schema.py | 8 +++-- .../api/schema_manage/edge_label.py | 4 ++- .../src/pyhugegraph/api/services.py | 3 +- .../src/pyhugegraph/api/traverser.py | 12 +++++-- .../pyhugegraph/example/hugegraph_example.py | 8 +++-- .../src/pyhugegraph/example/hugegraph_test.py | 7 ++-- .../structure/vertex_label_data.py | 5 ++- .../src/pyhugegraph/utils/huge_router.py | 4 ++- .../src/pyhugegraph/utils/util.py | 6 ++-- .../src/tests/api/test_traverser.py | 32 ++++++++++++++----- .../src/tests/client_utils.py | 32 ++++++++++++------- 42 files changed, 208 insertions(+), 148 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py index 96db7da0a..788c62082 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py @@ -32,7 +32,9 @@ def admin_http_api(router: APIRouter, log_stream): async def log_stream_api(req: LogStreamRequest): if admin_settings.admin_token != req.admin_token: raise generate_response( - RAGResponse(status_code=status.HTTP_403_FORBIDDEN, message="Invalid admin_token") # pylint: disable=E0702 + RAGResponse( + status_code=status.HTTP_403_FORBIDDEN, message="Invalid admin_token" + ) # pylint: disable=E0702 ) log_path = os.path.join("logs", req.log_file) diff --git a/hugegraph-llm/src/hugegraph_llm/config/generate.py b/hugegraph-llm/src/hugegraph_llm/config/generate.py index 9574b7b06..1bd7adea8 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/generate.py +++ b/hugegraph-llm/src/hugegraph_llm/config/generate.py @@ -18,7 +18,13 @@ import argparse -from hugegraph_llm.config import PromptConfig, admin_settings, huge_settings, index_settings, llm_settings +from hugegraph_llm.config import ( + PromptConfig, + admin_settings, + huge_settings, + index_settings, + llm_settings, +) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Generate hugegraph-llm config file") diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 7eecefb9e..0b48a84dc 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -108,7 +108,9 @@ def create_admin_block(): ) # Error message box, initially hidden - error_message = gr.Textbox(label="", visible=False, interactive=False, elem_classes="error-message") + error_message = gr.Textbox( + label="", visible=False, interactive=False, elem_classes="error-message" + ) # Button to submit password submit_button = gr.Button("Submit") diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 6978158c8..a6f57fc67 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -104,6 +104,8 @@ def test_api_connection( except (json.decoder.JSONDecodeError, AttributeError) as e: raise gr.Error(resp.text) from e return resp.status_code + + def apply_vector_engine(engine: str): # Persist the vector engine selection setattr(index_settings, "cur_vector_index", engine) @@ -125,18 +127,22 @@ def apply_vector_engine_backend( ) -> int: """Test connection and persist per-engine connection settings""" status_code = -1 - + # Test connection first try: if engine == "Milvus": from pymilvus import connections, utility - connections.connect(host=host, port=int(port or 19530), user=user or "", password=password or "") + + connections.connect( + host=host, port=int(port or 19530), user=user or "", password=password or "" + ) # Test if we can list collections _ = utility.list_collections() connections.disconnect("default") status_code = 200 elif engine == "Qdrant": from qdrant_client import QdrantClient + client = QdrantClient(host=host, port=int(port or 6333), api_key=api_key) # Test if we can get collections _ = client.get_collections() @@ -152,7 +158,7 @@ def apply_vector_engine_backend( if origin_call is None: raise gr.Error(msg) from e return -1 - + # Persist settings after successful test if engine == "Milvus": if host is not None: @@ -748,7 +754,9 @@ def vector_engine_settings(engine): gr.Textbox(value=index_settings.milvus_host, label="host"), gr.Textbox(value=str(index_settings.milvus_port), label="port"), gr.Textbox(value=index_settings.milvus_user, label="user"), - gr.Textbox(value=index_settings.milvus_password, label="password", type="password"), + gr.Textbox( + value=index_settings.milvus_password, label="password", type="password" + ), ] apply_backend_button = gr.Button("Apply Configuration") apply_backend_button.click( @@ -759,7 +767,11 @@ def vector_engine_settings(engine): qdrant_inputs = [ gr.Textbox(value=index_settings.qdrant_host, label="host"), gr.Textbox(value=str(index_settings.qdrant_port), label="port"), - gr.Textbox(value=(index_settings.qdrant_api_key or ""), label="api_key", type="password"), + gr.Textbox( + value=(index_settings.qdrant_api_key or ""), + label="api_key", + type="password", + ), ] apply_backend_button = gr.Button("Apply Configuration") apply_backend_button.click( diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py index efa77602e..8b78328f3 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py @@ -58,7 +58,9 @@ def create_other_block(): async def lifespan(app: FastAPI): # pylint: disable=W0621 log.info("Starting background scheduler...") scheduler = AsyncIOScheduler() - scheduler.add_job(backup_data, trigger=CronTrigger(hour=1, minute=0), id="daily_backup", replace_existing=True) + scheduler.add_job( + backup_data, trigger=CronTrigger(hour=1, minute=0), id="daily_backup", replace_existing=True + ) scheduler.start() log.info("Starting vid embedding update task...") diff --git a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py index 4c4a42e3a..ee173b284 100644 --- a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py +++ b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py @@ -22,7 +22,9 @@ class ChunkSplitter: def __init__( - self, split_type: Literal["paragraph", "sentence"] = "paragraph", language: Literal["zh", "en"] = "zh" + self, + split_type: Literal["paragraph", "sentence"] = "paragraph", + language: Literal["zh", "en"] = "zh", ): if language == "zh": separators = ["\n\n", "\n", "。", "，", ""] diff --git a/hugegraph-llm/src/hugegraph_llm/flows/common.py b/hugegraph-llm/src/hugegraph_llm/flows/common.py index e2348466c..ae495b520 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/common.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/common.py @@ -46,9 +46,7 @@ def post_deal(self, *args, **kwargs): """ pass - async def post_deal_stream( - self, pipeline=None - ) -> AsyncGenerator[Dict[str, Any], None]: + async def post_deal_stream(self, pipeline=None) -> AsyncGenerator[Dict[str, Any], None]: """ Streaming post-processing interface. Subclasses can override this method as needed. diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py index 5feb3d471..8cdd6da4c 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py @@ -69,12 +69,8 @@ def prepare( prepared_input.graph_vector_answer = graph_vector_answer prepared_input.gremlin_tmpl_num = gremlin_tmpl_num prepared_input.gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt - prepared_input.max_graph_items = ( - max_graph_items or huge_settings.max_graph_items - ) - prepared_input.topk_per_keyword = ( - topk_per_keyword or huge_settings.topk_per_keyword - ) + prepared_input.max_graph_items = max_graph_items or huge_settings.max_graph_items + prepared_input.topk_per_keyword = topk_per_keyword or huge_settings.topk_per_keyword prepared_input.topk_return_results = ( topk_return_results or huge_settings.topk_return_results ) @@ -123,18 +119,14 @@ def build_flow(self, **kwargs): {only_schema_node, only_semantic_id_query_node}, "only_graph", ) - pipeline.registerGElement( - merge_rerank_node, {only_graph_query_node}, "merge_one" - ) + pipeline.registerGElement(merge_rerank_node, {only_graph_query_node}, "merge_one") pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "graph") log.info("RAGGraphOnlyFlow pipeline built successfully") return pipeline def post_deal(self, pipeline=None): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) + return json.dumps({"error": "No pipeline provided"}, ensure_ascii=False, indent=2) try: res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() log.info("RAGGraphOnlyFlow post processing success") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py index 2f4a2bfa2..0e99f0252 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py @@ -71,15 +71,11 @@ def prepare( prepared_input.graph_ratio = graph_ratio prepared_input.gremlin_tmpl_num = gremlin_tmpl_num prepared_input.gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt - prepared_input.max_graph_items = ( - max_graph_items or huge_settings.max_graph_items - ) + prepared_input.max_graph_items = max_graph_items or huge_settings.max_graph_items prepared_input.topk_return_results = ( topk_return_results or huge_settings.topk_return_results ) - prepared_input.topk_per_keyword = ( - topk_per_keyword or huge_settings.topk_per_keyword - ) + prepared_input.topk_per_keyword = topk_per_keyword or huge_settings.topk_per_keyword prepared_input.vector_dis_threshold = ( vector_dis_threshold or huge_settings.vector_dis_threshold ) @@ -119,27 +115,17 @@ def build_flow(self, **kwargs): # Register nodes and their dependencies pipeline.registerGElement(vector_query_node, set(), "vector") pipeline.registerGElement(keyword_extract_node, set(), "keyword") - pipeline.registerGElement( - semantic_id_query_node, {keyword_extract_node}, "semantic" - ) + pipeline.registerGElement(semantic_id_query_node, {keyword_extract_node}, "semantic") pipeline.registerGElement(schema_node, set(), "schema") - pipeline.registerGElement( - graph_query_node, {schema_node, semantic_id_query_node}, "graph" - ) - pipeline.registerGElement( - merge_rerank_node, {graph_query_node, vector_query_node}, "merge" - ) - pipeline.registerGElement( - answer_synthesize_node, {merge_rerank_node}, "graph_vector" - ) + pipeline.registerGElement(graph_query_node, {schema_node, semantic_id_query_node}, "graph") + pipeline.registerGElement(merge_rerank_node, {graph_query_node, vector_query_node}, "merge") + pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "graph_vector") log.info("RAGGraphVectorFlow pipeline built successfully") return pipeline def post_deal(self, pipeline=None): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) + return json.dumps({"error": "No pipeline provided"}, ensure_ascii=False, indent=2) try: res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() log.info("RAGGraphVectorFlow post processing success") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py index f62e574bb..93dd25a4a 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py @@ -78,9 +78,7 @@ def build_flow(self, **kwargs): def post_deal(self, pipeline=None): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) + return json.dumps({"error": "No pipeline provided"}, ensure_ascii=False, indent=2) try: res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() log.info("RAGRawFlow post processing success") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py index c727eacce..e8df945ae 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py @@ -93,18 +93,14 @@ def build_flow(self, **kwargs): # Register nodes and dependencies, keep naming consistent with original pipeline.registerGElement(only_vector_query_node, set(), "only_vector") - pipeline.registerGElement( - merge_rerank_node, {only_vector_query_node}, "merge_two" - ) + pipeline.registerGElement(merge_rerank_node, {only_vector_query_node}, "merge_two") pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "vector") log.info("RAGVectorOnlyFlow pipeline built successfully") return pipeline def post_deal(self, pipeline=None): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) + return json.dumps({"error": "No pipeline provided"}, ensure_ascii=False, indent=2) try: res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() log.info("RAGVectorOnlyFlow post processing success") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index 5afa1bf8e..0fd550e20 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -158,9 +158,7 @@ async def schedule_stream_flow(self, flow: str, *args, **kwargs): else: try: # fetch pipeline & prepare input for flow - prepared_input: WkFlowInput = pipeline.getGParamWithNoEmpty( - "wkflow_input" - ) + prepared_input: WkFlowInput = pipeline.getGParamWithNoEmpty("wkflow_input") prepared_input.stream = True flow.prepare(prepared_input, *args, **kwargs) status = pipeline.run() diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index ca7761ecd..4342c90ee 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -185,7 +185,9 @@ def get_vector_index_info(self) -> Dict: def clean(*name: str): name_str = "_".join(name) client = QdrantClient( - host=index_settings.qdrant_host, port=index_settings.qdrant_port, api_key=index_settings.qdrant_api_key + host=index_settings.qdrant_host, + port=index_settings.qdrant_port, + api_key=index_settings.qdrant_api_key, ) collections = client.get_collections().collections collection_names = [collection.name for collection in collections] @@ -209,7 +211,9 @@ def from_name(embed_dim: int, *name: str) -> "QdrantVectorIndex": def exist(*name: str) -> bool: name_str = "_".join(name) client = QdrantClient( - host=index_settings.qdrant_host, port=index_settings.qdrant_port, api_key=index_settings.qdrant_api_key + host=index_settings.qdrant_host, + port=index_settings.qdrant_port, + api_key=index_settings.qdrant_api_key, ) collections = client.get_collections().collections collection_names = [collection.name for collection in collections] diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py index ae65ccb33..564e0eb2c 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py @@ -38,16 +38,12 @@ def node_init(self): return CStatus(-1, "graph_name is required in wk_input") max_deep = self.wk_input.max_deep or 2 - max_graph_items = ( - self.wk_input.max_graph_items or huge_settings.max_graph_items - ) + max_graph_items = self.wk_input.max_graph_items or huge_settings.max_graph_items max_v_prop_len = self.wk_input.max_v_prop_len or 2048 max_e_prop_len = self.wk_input.max_e_prop_len or 256 prop_to_match = self.wk_input.prop_to_match num_gremlin_generate_example = self.wk_input.gremlin_tmpl_num or -1 - gremlin_prompt = ( - self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt - ) + gremlin_prompt = self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt # Initialize GraphRAGQuery operator self.graph_rag_query = GraphRAGQuery( diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py index 88b019e33..fc09b0ec6 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -38,9 +38,7 @@ def node_init(self): example_num = 2 # Clamp to [0, 10] example_num = max(0, min(10, example_num)) - self.operator = GremlinExampleIndexQuery( - embedding=embedding, num_examples=example_num - ) + self.operator = GremlinExampleIndexQuery(embedding=embedding, num_examples=example_num) return CStatus() def operator_schedule(self, data_json: Dict[str, Any]): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py index db535bd6f..2897ba181 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -40,9 +40,7 @@ def node_init(self): embedding = Embeddings().get_embedding() by = self.wk_input.semantic_by or "keywords" - topk_per_keyword = ( - self.wk_input.topk_per_keyword or huge_settings.topk_per_keyword - ) + topk_per_keyword = self.wk_input.topk_per_keyword or huge_settings.topk_per_keyword topk_per_query = self.wk_input.topk_per_query or 10 vector_dis_threshold = ( self.wk_input.vector_dis_threshold or huge_settings.vector_dis_threshold @@ -80,9 +78,7 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: semantic_result = self.semantic_id_query.run(data_json) match_vids = semantic_result.get("match_vids", []) - log.info( - f"Semantic query completed, found {len(match_vids)} matching vertex IDs" - ) + log.info(f"Semantic query completed, found {len(match_vids)} matching vertex IDs") return semantic_result diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py index d4c5c8026..c9eb8c4e9 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py @@ -35,9 +35,7 @@ def node_init(self): try: # 从 wk_input 中读取用户配置参数 embedding = Embeddings().get_embedding() - max_items = ( - self.wk_input.max_items if self.wk_input.max_items is not None else 3 - ) + max_items = self.wk_input.max_items if self.wk_input.max_items is not None else 3 self.operator = VectorIndexQuery(embedding=embedding, topk=max_items) return super().node_init() diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py index 22b970b4a..b32bd13e8 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py @@ -75,9 +75,7 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: if result.get("graph_vector_answer"): answer_types.append("graph_vector") - log.info( - f"Answer synthesis completed for types: {', '.join(answer_types)}" - ) + log.info(f"Answer synthesis completed for types: {', '.join(answer_types)}") # Print enabled answer types according to self.wk_input configuration wk_input_types = [] diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py index 76fc06eb3..ea77c6580 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py @@ -34,15 +34,9 @@ def node_init(self): """ try: max_keywords = ( - self.wk_input.max_keywords - if self.wk_input.max_keywords is not None - else 5 - ) - language = ( - self.wk_input.language - if self.wk_input.language is not None - else "english" + self.wk_input.max_keywords if self.wk_input.max_keywords is not None else 5 ) + language = self.wk_input.language if self.wk_input.language is not None else "english" extract_template = self.wk_input.keywords_extract_prompt self.operator = KeywordExtract( diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py index a36831526..bcb995f82 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py @@ -28,9 +28,7 @@ def _stable_schema_string(state_json: Dict[str, Any]) -> str: if "simple_schema" in state_json and state_json["simple_schema"] is not None: - return json.dumps( - state_json["simple_schema"], ensure_ascii=False, sort_keys=True - ) + return json.dumps(state_json["simple_schema"], ensure_ascii=False, sort_keys=True) if "schema" in state_json and state_json["schema"] is not None: return json.dumps(state_json["schema"], ensure_ascii=False, sort_keys=True) return "" diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index d052f2f6e..fb385a59c 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -26,7 +26,12 @@ # FIXME: we need keep the logic same with build_semantic_index.py class BuildGremlinExampleIndex: - def __init__(self, embedding: BaseEmbedding, examples: List[Dict[str, str]], vector_index: type[VectorStoreBase]): + def __init__( + self, + embedding: BaseEmbedding, + examples: List[Dict[str, str]], + vector_index: type[VectorStoreBase], + ): self.vector_index_name = "gremlin_examples" self.examples = examples self.embedding = embedding diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index 3f660880d..e49ca04ee 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -36,7 +36,10 @@ class GremlinExampleIndexQuery: def __init__( - self, vector_index: type[VectorStoreBase], embedding: Optional[BaseEmbedding] = None, num_examples: int = 1 + self, + vector_index: type[VectorStoreBase], + embedding: Optional[BaseEmbedding] = None, + num_examples: int = 1, ): self.embedding = embedding or Embeddings().get_embedding() self.num_examples = num_examples @@ -70,7 +73,9 @@ def _build_default_example_index(self): with ThreadPoolExecutor() as executor: embeddings = list( tqdm( - executor.map(self.embedding.get_text_embedding, [row["query"] for row in properties]), + executor.map( + self.embedding.get_text_embedding, [row["query"] for row in properties] + ), total=len(properties), ) ) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py index 0ae1dd4f6..5ced65b41 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py @@ -25,7 +25,9 @@ class VectorIndexQuery: - def __init__(self, vector_index: type[VectorStoreBase], embedding: BaseEmbedding, topk: int = 3): + def __init__( + self, vector_index: type[VectorStoreBase], embedding: BaseEmbedding, topk: int = 3 + ): self.embedding = embedding self.topk = topk self.vector_index = vector_index.from_name( diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py index 5913ea307..819ef25fe 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py @@ -51,7 +51,9 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: llm_output = self.llm.generate(prompt=prompt) data["triples"] = [] extract_triples_by_regex(llm_output, data) - print(f"LLM {self.__class__.__name__} input:{prompt} \n output: {llm_output} \n data: {data}") + print( + f"LLM {self.__class__.__name__} input:{prompt} \n output: {llm_output} \n data: {data}" + ) data["call_count"] = data.get("call_count", 0) + 1 return data diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py index 972c828ec..928948413 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py @@ -86,7 +86,10 @@ def _extract_schema(self, response: str) -> Dict[str, Any]: raise RuntimeError("Invalid JSON response from LLM") from e def build_prompt( - self, raw_texts: List[str], query_examples: List[Dict[str, str]], few_shot_schema: Dict[str, Any] + self, + raw_texts: List[str], + query_examples: List[Dict[str, str]], + few_shot_schema: Dict[str, Any], ) -> str: return self.schema_prompt.format( raw_texts=self._format_raw_texts(raw_texts), diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py index 6beeb0291..38eabb16e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/unstructured_data_utils.py @@ -20,7 +20,10 @@ import re REGEX = ( - r"Nodes:\s+(.*?)\s?\s?" r"Relationships:\s?\s?" r"NodesSchemas:\s+(.*?)\s?\s?" r"RelationshipsSchemas:\s?\s?(.*)" + r"Nodes:\s+(.*?)\s?\s?" + r"Relationships:\s?\s?" + r"NodesSchemas:\s+(.*?)\s?\s?" + r"RelationshipsSchemas:\s?\s?(.*)" ) INTERNAL_REGEX = r"\[(.*?)\]" JSON_REGEX = r"\{.*\}" diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 3a6fd3c1c..51ff0b306 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -198,11 +198,7 @@ def to_json(self): dict: A dictionary containing non-None instance members and their serialized values. """ # Only export instance attributes (excluding methods and class attributes) whose values are not None - return { - k: v - for k, v in self.__dict__.items() - if not k.startswith("_") and v is not None - } + return {k: v for k, v in self.__dict__.items() if not k.startswith("_") and v is not None} # Implement a method that assigns keys from data_json as WkFlowState member variables def assign_from_json(self, data_json: dict): diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index 6965b6239..147c0074c 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -140,7 +140,10 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): elif filename == "vertices.json": data_full = client.gremlin().exec(query)["data"][0]["vertices"] data = ( - [{key: value for key, value in vertex.items() if key != "id"} for vertex in data_full] + [ + {key: value for key, value in vertex.items() if key != "id"} + for vertex in data_full + ] if all_pk_flag else data_full ) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/log.py b/hugegraph-llm/src/hugegraph_llm/utils/log.py index 3d956a51b..b64017454 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/log.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/log.py @@ -27,7 +27,11 @@ # Initialize the root logger first with Rich handler root_logger = init_logger( - log_output=LOG_FILE, log_level=INFO, logger_name="root", propagate_logs=True, stdout_logging=True + log_output=LOG_FILE, + log_level=INFO, + logger_name="root", + propagate_logs=True, + stdout_logging=True, ) # Initialize custom logger diff --git a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py index 734d87263..7ad914468 100644 --- a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py +++ b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py @@ -32,4 +32,6 @@ def test_stream_generate(self): def on_token_callback(chunk): print(chunk, end="", flush=True) - ollama_client.generate_streaming(prompt="What is the capital of France?", on_token_callback=on_token_callback) + ollama_client.generate_streaming( + prompt="What is the capital of France?", on_token_callback=on_token_callback + ) diff --git a/hugegraph-python-client/src/pyhugegraph/api/auth.py b/hugegraph-python-client/src/pyhugegraph/api/auth.py index ea3695b99..d127c4f6d 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/auth.py +++ b/hugegraph-python-client/src/pyhugegraph/api/auth.py @@ -31,7 +31,9 @@ def list_users(self, limit=None): return self._invoke_request(params=params) @router.http("POST", "auth/users") - def create_user(self, user_name, user_password, user_phone=None, user_email=None) -> Optional[Dict]: + def create_user( + self, user_name, user_password, user_phone=None, user_email=None + ) -> Optional[Dict]: return self._invoke_request( data=json.dumps( { @@ -116,7 +118,9 @@ def revoke_accesses(self, access_id) -> Optional[Dict]: # pylint: disable=unuse return self._invoke_request() @router.http("PUT", "auth/accesses/{access_id}") - def modify_accesses(self, access_id, access_description) -> Optional[Dict]: # pylint: disable=unused-argument + def modify_accesses( + self, access_id, access_description + ) -> Optional[Dict]: # pylint: disable=unused-argument # The permission of access can\'t be updated data = {"access_description": access_description} return self._invoke_request(data=json.dumps(data)) @@ -130,7 +134,9 @@ def list_accesses(self) -> Optional[Dict]: return self._invoke_request() @router.http("POST", "auth/targets") - def create_target(self, target_name, target_graph, target_url, target_resources) -> Optional[Dict]: + def create_target( + self, target_name, target_graph, target_url, target_resources + ) -> Optional[Dict]: return self._invoke_request( data=json.dumps( { @@ -167,7 +173,9 @@ def update_target( ) @router.http("GET", "auth/targets/{target_id}") - def get_target(self, target_id, response=None) -> Optional[Dict]: # pylint: disable=unused-argument + def get_target( + self, target_id, response=None + ) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "auth/targets") @@ -184,7 +192,9 @@ def delete_belong(self, belong_id) -> None: # pylint: disable=unused-argument return self._invoke_request() @router.http("PUT", "auth/belongs/{belong_id}") - def update_belong(self, belong_id, description) -> Optional[Dict]: # pylint: disable=unused-argument + def update_belong( + self, belong_id, description + ) -> Optional[Dict]: # pylint: disable=unused-argument data = {"belong_description": description} return self._invoke_request(data=json.dumps(data)) diff --git a/hugegraph-python-client/src/pyhugegraph/api/schema.py b/hugegraph-python-client/src/pyhugegraph/api/schema.py index 3576ce66b..8095887b0 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/schema.py +++ b/hugegraph-python-client/src/pyhugegraph/api/schema.py @@ -68,7 +68,9 @@ def getSchema(self, _format: str = "json") -> Optional[Dict]: # pylint: disable return self._invoke_request() @router.http("GET", "schema/propertykeys/{property_name}") - def getPropertyKey(self, property_name) -> Optional[PropertyKeyData]: # pylint: disable=unused-argument + def getPropertyKey( + self, property_name + ) -> Optional[PropertyKeyData]: # pylint: disable=unused-argument if response := self._invoke_request(): return PropertyKeyData(response) return None @@ -93,7 +95,9 @@ def getVertexLabels(self) -> Optional[List[VertexLabelData]]: return None @router.http("GET", "schema/edgelabels/{label_name}") - def getEdgeLabel(self, label_name: str) -> Optional[EdgeLabelData]: # pylint: disable=unused-argument + def getEdgeLabel( + self, label_name: str + ) -> Optional[EdgeLabelData]: # pylint: disable=unused-argument if response := self._invoke_request(): return EdgeLabelData(response) log.error("EdgeLabel not found: %s", str(response)) diff --git a/hugegraph-python-client/src/pyhugegraph/api/schema_manage/edge_label.py b/hugegraph-python-client/src/pyhugegraph/api/schema_manage/edge_label.py index 91608fe6a..93f218001 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/schema_manage/edge_label.py +++ b/hugegraph-python-client/src/pyhugegraph/api/schema_manage/edge_label.py @@ -150,7 +150,9 @@ def append(self): def eliminate(self): name = self._parameter_holder.get_value("name") user_data = ( - self._parameter_holder.get_value("user_data") if self._parameter_holder.get_value("user_data") else {} + self._parameter_holder.get_value("user_data") + if self._parameter_holder.get_value("user_data") + else {} ) path = f"schema/edgelabels/{name}?action=eliminate" data = {"name": name, "user_data": user_data} diff --git a/hugegraph-python-client/src/pyhugegraph/api/services.py b/hugegraph-python-client/src/pyhugegraph/api/services.py index 4fac4aa69..f353673db 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/services.py +++ b/hugegraph-python-client/src/pyhugegraph/api/services.py @@ -125,6 +125,7 @@ def delete_service(self, graphspace: str, service: str): # pylint: disable=unus None """ return self._sess.request( - f"/graphspaces/{graphspace}/services/{service}" f"?confirm_message=I'm sure to delete the service", + f"/graphspaces/{graphspace}/services/{service}" + f"?confirm_message=I'm sure to delete the service", "DELETE", ) diff --git a/hugegraph-python-client/src/pyhugegraph/api/traverser.py b/hugegraph-python-client/src/pyhugegraph/api/traverser.py index 199ed0167..2f226522b 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/traverser.py +++ b/hugegraph-python-client/src/pyhugegraph/api/traverser.py @@ -49,7 +49,9 @@ def shortest_path(self, source_id, target_id, max_depth): # pylint: disable=unu "GET", 'traversers/allshortestpaths?source="{source_id}"&target="{target_id}"&max_depth={max_depth}', ) - def all_shortest_paths(self, source_id, target_id, max_depth): # pylint: disable=unused-argument + def all_shortest_paths( + self, source_id, target_id, max_depth + ): # pylint: disable=unused-argument return self._invoke_request() @router.http( @@ -57,7 +59,9 @@ def all_shortest_paths(self, source_id, target_id, max_depth): # pylint: disabl 'traversers/weightedshortestpath?source="{source_id}"&target="{target_id}"' "&weight={weight}&max_depth={max_depth}", ) - def weighted_shortest_path(self, source_id, target_id, weight, max_depth): # pylint: disable=unused-argument + def weighted_shortest_path( + self, source_id, target_id, weight, max_depth + ): # pylint: disable=unused-argument return self._invoke_request() @router.http( @@ -126,7 +130,9 @@ def advanced_paths( ) @router.http("POST", "traversers/customizedpaths") - def customized_paths(self, sources, steps, sort_by="INCR", with_vertex=True, capacity=-1, limit=-1): + def customized_paths( + self, sources, steps, sort_by="INCR", with_vertex=True, capacity=-1, limit=-1 + ): return self._invoke_request( data=json.dumps( { diff --git a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py index de2bb67ee..d5cc0eb9d 100644 --- a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py +++ b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py @@ -18,7 +18,9 @@ from pyhugegraph.client import PyHugeClient if __name__ == "__main__": - client = PyHugeClient(url="http://127.0.0.1:8080", user="admin", pwd="admin", graph="hugegraph", graphspace=None) + client = PyHugeClient( + url="http://127.0.0.1:8080", user="admin", pwd="admin", graph="hugegraph", graphspace=None + ) """schema""" schema = client.schema() @@ -27,7 +29,9 @@ schema.vertexLabel("Person").properties("name", "birthDate").usePrimaryKeyId().primaryKeys( "name" ).ifNotExist().create() - schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys("name").ifNotExist().create() + schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys( + "name" + ).ifNotExist().create() schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() print(schema.getVertexLabels()) diff --git a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_test.py b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_test.py index 075d069b7..2bfe6ea97 100644 --- a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_test.py +++ b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_test.py @@ -31,14 +31,17 @@ def __init__( from pyhugegraph.client import PyHugeClient except ImportError: raise ValueError( - "Please install HugeGraph Python client first: " "`pip3 install hugegraph-python-client`" + "Please install HugeGraph Python client first: " + "`pip3 install hugegraph-python-client`" ) from ImportError self.username = username self.password = password self.url = url self.graph = graph - self.client = PyHugeClient(url=url, user=username, pwd=password, graph=graph, graphspace=None) + self.client = PyHugeClient( + url=url, user=username, pwd=password, graph=graph, graphspace=None + ) self.schema = "" def exec(self, query) -> str: diff --git a/hugegraph-python-client/src/pyhugegraph/structure/vertex_label_data.py b/hugegraph-python-client/src/pyhugegraph/structure/vertex_label_data.py index aaee1370f..39da4eebb 100644 --- a/hugegraph-python-client/src/pyhugegraph/structure/vertex_label_data.py +++ b/hugegraph-python-client/src/pyhugegraph/structure/vertex_label_data.py @@ -65,5 +65,8 @@ def enableLabelIndex(self): return self.__enable_label_index def __repr__(self): - res = f"name: {self.__name}, primary_keys: {self.__primary_keys}, " f"properties: {self.__properties}" + res = ( + f"name: {self.__name}, primary_keys: {self.__primary_keys}, " + f"properties: {self.__properties}" + ) return res diff --git a/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py b/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py index 7acf2fdfa..f4a38a418 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py @@ -145,7 +145,9 @@ def wrapper(self: "HGraphContext", *args: Any, **kwargs: Any) -> Any: class RouterMixin: - def _invoke_request_registered(self, placeholders: dict = None, validator=ResponseValidation(), **kwargs: Any): + def _invoke_request_registered( + self, placeholders: dict = None, validator=ResponseValidation(), **kwargs: Any + ): """ Make an HTTP request using the stored partial request function. Args: diff --git a/hugegraph-python-client/src/pyhugegraph/utils/util.py b/hugegraph-python-client/src/pyhugegraph/utils/util.py index 3fe300092..76770a818 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/util.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/util.py @@ -34,7 +34,8 @@ def create_exception(response_content): data = json.loads(response_content) if "ServiceUnavailableException" in data.get("exception", ""): raise ServiceUnavailableException( - f'ServiceUnavailableException, "message": "{data["message"]}",' f' "cause": "{data["cause"]}"' + f'ServiceUnavailableException, "message": "{data["message"]}",' + f' "cause": "{data["cause"]}"' ) except (json.JSONDecodeError, KeyError) as e: raise Exception(f"Error parsing response content: {response_content}") from e @@ -57,8 +58,7 @@ def check_if_success(response, error=None): req = response.request req_body = req.body if req.body else "Empty body" response_body = response.text if response.text else "Empty body" - log.error( - ) + log.error() raise error return True diff --git a/hugegraph-python-client/src/tests/api/test_traverser.py b/hugegraph-python-client/src/tests/api/test_traverser.py index ae44cf6f8..70c206acc 100644 --- a/hugegraph-python-client/src/tests/api/test_traverser.py +++ b/hugegraph-python-client/src/tests/api/test_traverser.py @@ -55,7 +55,9 @@ def test_traverser_operations(self): self.assertEqual(k_out_result["vertices"], ["1:peter", "2:ripple"]) k_neighbor_result = self.traverser.k_neighbor(marko, 2) - self.assertEqual(k_neighbor_result["vertices"], ["1:peter", "1:josh", "2:lop", "2:ripple", "1:vadas"]) + self.assertEqual( + k_neighbor_result["vertices"], ["1:peter", "1:josh", "2:lop", "2:ripple", "1:vadas"] + ) same_neighbors_result = self.traverser.same_neighbors(marko, josh) self.assertEqual(same_neighbors_result["same_neighbors"], ["2:lop"]) @@ -67,10 +69,16 @@ def test_traverser_operations(self): self.assertEqual(shortest_path_result["path"], ["1:marko", "1:josh", "2:ripple"]) all_shortest_paths_result = self.traverser.all_shortest_paths(marko, ripple, 3) - self.assertEqual(all_shortest_paths_result["paths"], [{"objects": ["1:marko", "1:josh", "2:ripple"]}]) + self.assertEqual( + all_shortest_paths_result["paths"], [{"objects": ["1:marko", "1:josh", "2:ripple"]}] + ) - weighted_shortest_path_result = self.traverser.weighted_shortest_path(marko, ripple, "weight", 3) - self.assertEqual(weighted_shortest_path_result["vertices"], ["1:marko", "1:josh", "2:ripple"]) + weighted_shortest_path_result = self.traverser.weighted_shortest_path( + marko, ripple, "weight", 3 + ) + self.assertEqual( + weighted_shortest_path_result["vertices"], ["1:marko", "1:josh", "2:ripple"] + ) single_source_shortest_path_result = self.traverser.single_source_shortest_path(marko, 2) self.assertEqual( @@ -84,7 +92,9 @@ def test_traverser_operations(self): }, ) - multi_node_shortest_path_result = self.traverser.multi_node_shortest_path([marko, josh], max_depth=2) + multi_node_shortest_path_result = self.traverser.multi_node_shortest_path( + [marko, josh], max_depth=2 + ) self.assertEqual( multi_node_shortest_path_result["vertices"], [ @@ -121,7 +131,9 @@ def test_traverser_operations(self): } ], ) - self.assertEqual(customized_paths_result["paths"], [{"objects": ["1:marko", "2:lop"], "weights": [8.0]}]) + self.assertEqual( + customized_paths_result["paths"], [{"objects": ["1:marko", "2:lop"], "weights": [8.0]}] + ) sources = {"ids": [], "label": "person", "properties": {"name": "vadas"}} @@ -174,11 +186,15 @@ def test_traverser_operations(self): sources = {"ids": ["2:lop", "2:ripple"]} path_patterns = [{"steps": [{"direction": "IN", "labels": ["created"], "max_degree": -1}]}] - customized_crosspoints_result = self.traverser.customized_crosspoints(sources, path_patterns) + customized_crosspoints_result = self.traverser.customized_crosspoints( + sources, path_patterns + ) self.assertEqual(customized_crosspoints_result["crosspoints"], ["1:josh"]) rings_result = self.traverser.rings(marko, 3) - self.assertEqual(rings_result["rings"], [{"objects": ["1:marko", "2:lop", "1:josh", "1:marko"]}]) + self.assertEqual( + rings_result["rings"], [{"objects": ["1:marko", "2:lop", "1:josh", "1:marko"]}] + ) rays_result = self.traverser.rays(marko, 2) self.assertEqual( diff --git a/hugegraph-python-client/src/tests/client_utils.py b/hugegraph-python-client/src/tests/client_utils.py index 11cbb4a55..f711072b8 100644 --- a/hugegraph-python-client/src/tests/client_utils.py +++ b/hugegraph-python-client/src/tests/client_utils.py @@ -59,21 +59,23 @@ def init_property_key(self): def init_vertex_label(self): schema = self.schema - schema.vertexLabel("person").properties("name", "age", "city").primaryKeys("name").nullableKeys( - "city" - ).ifNotExist().create() - schema.vertexLabel("software").properties("name", "lang", "price").primaryKeys("name").nullableKeys( - "price" - ).ifNotExist().create() + schema.vertexLabel("person").properties("name", "age", "city").primaryKeys( + "name" + ).nullableKeys("city").ifNotExist().create() + schema.vertexLabel("software").properties("name", "lang", "price").primaryKeys( + "name" + ).nullableKeys("price").ifNotExist().create() schema.vertexLabel("book").useCustomizeStringId().properties("name", "price").nullableKeys( "price" ).ifNotExist().create() def init_edge_label(self): schema = self.schema - schema.edgeLabel("knows").sourceLabel("person").targetLabel("person").multiTimes().properties( - "date", "city" - ).sortKeys("date").nullableKeys("city").ifNotExist().create() + schema.edgeLabel("knows").sourceLabel("person").targetLabel( + "person" + ).multiTimes().properties("date", "city").sortKeys("date").nullableKeys( + "city" + ).ifNotExist().create() schema.edgeLabel("created").sourceLabel("person").targetLabel("software").properties( "date", "city" ).nullableKeys("city").ifNotExist().create() @@ -82,10 +84,16 @@ def init_index_label(self): schema = self.schema schema.indexLabel("personByCity").onV("person").by("city").secondary().ifNotExist().create() schema.indexLabel("personByAge").onV("person").by("age").range().ifNotExist().create() - schema.indexLabel("softwareByPrice").onV("software").by("price").range().ifNotExist().create() - schema.indexLabel("softwareByLang").onV("software").by("lang").secondary().ifNotExist().create() + schema.indexLabel("softwareByPrice").onV("software").by( + "price" + ).range().ifNotExist().create() + schema.indexLabel("softwareByLang").onV("software").by( + "lang" + ).secondary().ifNotExist().create() schema.indexLabel("knowsByDate").onE("knows").by("date").secondary().ifNotExist().create() - schema.indexLabel("createdByDate").onE("created").by("date").secondary().ifNotExist().create() + schema.indexLabel("createdByDate").onE("created").by( + "date" + ).secondary().ifNotExist().create() def init_vertices(self): graph = self.graph From cb657600b17232fcfc8bc845f9414b742acd0f94 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Sun, 12 Oct 2025 14:00:31 +0800 Subject: [PATCH 41/71] fix pylint --- hugegraph-llm/src/hugegraph_llm/flows/common.py | 4 ++-- .../hugegraph_llm/flows/rag_flow_graph_only.py | 2 +- .../flows/rag_flow_graph_vector.py | 2 +- .../src/hugegraph_llm/flows/rag_flow_raw.py | 2 +- .../hugegraph_llm/flows/rag_flow_vector_only.py | 2 +- .../hugegraph_llm/models/embeddings/litellm.py | 17 ++++++++++++++++- .../hugegraph_llm/models/embeddings/ollama.py | 5 +++++ .../nodes/common_node/merge_rerank_node.py | 12 +++++++----- .../nodes/hugegraph_node/graph_query_node.py | 10 ++++++---- .../nodes/index_node/build_semantic_index.py | 1 - .../nodes/index_node/build_vector_index.py | 1 - .../index_node/gremlin_example_index_query.py | 1 - .../nodes/index_node/semantic_id_query_node.py | 11 ++++++----- .../nodes/index_node/vector_query_node.py | 5 ++--- .../nodes/llm_node/answer_synthesize_node.py | 6 +++--- .../nodes/llm_node/keyword_extract_node.py | 6 +++--- .../index_op/gremlin_example_index_query.py | 9 ++------- .../src/hugegraph_llm/state/ai_state.py | 4 ++-- 18 files changed, 58 insertions(+), 42 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/common.py b/hugegraph-llm/src/hugegraph_llm/flows/common.py index ae495b520..19a6ca2d4 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/common.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/common.py @@ -57,7 +57,7 @@ async def post_deal_stream(self, pipeline=None) -> AsyncGenerator[Dict[str, Any] return try: state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - log.info(f"{flow_name} post processing success") + log.info("%s post processing success", flow_name) stream_flow = state_json.get("stream_generator") if stream_flow is None: yield {"error": "No stream_generator found in workflow state"} @@ -65,5 +65,5 @@ async def post_deal_stream(self, pipeline=None) -> AsyncGenerator[Dict[str, Any] async for chunk in stream_flow: yield chunk except Exception as e: - log.error(f"{flow_name} post processing failed: {e}") + log.error("%s post processing failed: %s", flow_name, e) yield {"error": f"Post processing failed: {str(e)}"} diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py index 8cdd6da4c..e3397d639 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py @@ -137,7 +137,7 @@ def post_deal(self, pipeline=None): "graph_vector_answer": res.get("graph_vector_answer", ""), } except Exception as e: - log.error(f"RAGGraphOnlyFlow post processing failed: {e}") + log.error("RAGGraphOnlyFlow post processing failed: %s", e) return json.dumps( {"error": f"Post processing failed: {str(e)}"}, ensure_ascii=False, diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py index 0e99f0252..be9d8bbb9 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py @@ -136,7 +136,7 @@ def post_deal(self, pipeline=None): "graph_vector_answer": res.get("graph_vector_answer", ""), } except Exception as e: - log.error(f"RAGGraphVectorFlow post processing failed: {e}") + log.error("RAGGraphVectorFlow post processing failed: %s", e) return json.dumps( {"error": f"Post processing failed: {str(e)}"}, ensure_ascii=False, diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py index 93dd25a4a..d328d6c6e 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py @@ -89,7 +89,7 @@ def post_deal(self, pipeline=None): "graph_vector_answer": res.get("graph_vector_answer", ""), } except Exception as e: - log.error(f"RAGRawFlow post processing failed: {e}") + log.error("RAGRawFlow post processing failed: %s", e) return json.dumps( {"error": f"Post processing failed: {str(e)}"}, ensure_ascii=False, diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py index e8df945ae..0766903b1 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py @@ -111,7 +111,7 @@ def post_deal(self, pipeline=None): "graph_vector_answer": res.get("graph_vector_answer", ""), } except Exception as e: - log.error(f"RAGVectorOnlyFlow post processing failed: {e}") + log.error("RAGVectorOnlyFlow post processing failed: %s", e) return json.dumps( {"error": f"Post processing failed: {str(e)}"}, ensure_ascii=False, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py index 3eb47c566..c5effaacf 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py @@ -79,8 +79,23 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: log.error("Error in LiteLLM batch embedding call: %s", e) raise - async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: + async def async_get_text_embedding(self, text: str) -> List[float]: """Get embedding for a single text asynchronously.""" + try: + response = await aembedding( + model=self.model, + input=text, + api_key=self.api_key, + api_base=self.api_base, + ) + log.info("Token usage: %s", response.usage) + return response.data[0]["embedding"] + except (RateLimitError, APIConnectionError, APIError) as e: + log.error("Error in async LiteLLM embedding call: %s", e) + raise + + async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: + """Get embeddings for multiple texts asynchronously.""" try: response = await aembedding( model=self.model, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py index ac07836e1..c02590695 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py @@ -73,6 +73,11 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: response = self.client.embed(model=self.model, input=texts)["embeddings"] return [list(inner_sequence) for inner_sequence in response] + async def async_get_text_embedding(self, text: str) -> List[float]: + """Get embedding for a single text asynchronously.""" + response = await self.async_client.embeddings(model=self.model, prompt=text) + return list(response["embedding"]) + async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: # Ollama python client may not provide batch async embeddings; fallback per item results: List[List[float]] = [] diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py index 405be1bf0..d29ec2c4e 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py @@ -17,7 +17,7 @@ from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.common_op.merge_dedup_rerank import MergeDedupRerank from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.config import huge_settings from hugegraph_llm.utils.log import log @@ -53,7 +53,7 @@ def node_init(self): ) return super().node_init() except Exception as e: - log.error(f"Failed to initialize MergeRerankNode: {e}") + log.error("Failed to initialize MergeRerankNode: %s", e) from PyCGraph import CStatus return CStatus(-1, f"MergeRerankNode initialization failed: {e}") @@ -72,12 +72,14 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: merged_count = len(result.get("merged_result", [])) log.info( - f"Merge and rerank completed: {vector_count} vector results, " - f"{graph_count} graph results, {merged_count} merged results" + "Merge and rerank completed: %d vector results, %d graph results, %d merged results", + vector_count, + graph_count, + merged_count, ) return result except Exception as e: - log.error(f"Merge and rerank failed: {e}") + log.error("Merge and rerank failed: %s", e) return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py index 564e0eb2c..233e73e8c 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py @@ -13,8 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from typing import Dict, Any + +from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery from hugegraph_llm.config import huge_settings, prompt @@ -58,7 +59,7 @@ def node_init(self): return super().node_init() except Exception as e: - log.error(f"Failed to initialize GraphQueryNode: {e}") + log.error("Failed to initialize GraphQueryNode: %s", e) return CStatus(-1, f"GraphQueryNode initialization failed: {e}") @@ -79,11 +80,12 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: data_json.update(graph_result) log.info( - f"Graph query completed, found {len(data_json.get('graph_result', []))} results" + "Graph query completed, found %d results", + len(data_json.get("graph_result", [])), ) return data_json except Exception as e: - log.error(f"Graph query failed: {e}") + log.error("Graph query failed: %s", e) return 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 327e3cea8..6a4424f38 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 @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from hugegraph_llm.config import llm_settings from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py index a58a0dda4..28f2cb041 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from hugegraph_llm.config import llm_settings from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py index fc09b0ec6..8b9e0db4d 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -19,7 +19,6 @@ from PyCGraph import CStatus -from hugegraph_llm.config import llm_settings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.gremlin_example_index_query import ( GremlinExampleIndexQuery, diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py index 2897ba181..3e6df12ac 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -13,12 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from typing import Dict, Any + +from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.semantic_id_query import SemanticIdQuery from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.config import huge_settings from hugegraph_llm.utils.log import log @@ -57,7 +58,7 @@ def node_init(self): return super().node_init() except Exception as e: - log.error(f"Failed to initialize SemanticIdQueryNode: {e}") + log.error("Failed to initialize SemanticIdQueryNode: %s", e) return CStatus(-1, f"SemanticIdQueryNode initialization failed: {e}") @@ -78,10 +79,10 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: semantic_result = self.semantic_id_query.run(data_json) match_vids = semantic_result.get("match_vids", []) - log.info(f"Semantic query completed, found {len(match_vids)} matching vertex IDs") + log.info("Semantic query completed, found %d matching vertex IDs", len(match_vids)) return semantic_result except Exception as e: - log.error(f"Semantic query failed: {e}") + log.error("Semantic query failed: %s", e) return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py index c9eb8c4e9..f9af6c49d 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py @@ -14,7 +14,6 @@ # limitations under the License. from typing import Dict, Any -from hugegraph_llm.config import llm_settings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery from hugegraph_llm.models.embeddings.init_embedding import Embeddings @@ -40,7 +39,7 @@ def node_init(self): self.operator = VectorIndexQuery(embedding=embedding, topk=max_items) return super().node_init() except Exception as e: - log.error(f"Failed to initialize VectorQueryNode: {e}") + log.error("Failed to initialize VectorQueryNode: %s", e) from PyCGraph import CStatus return CStatus(-1, f"VectorQueryNode initialization failed: {e}") @@ -68,5 +67,5 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: return data_json except Exception as e: - log.error(f"Vector query failed: {e}") + log.error("Vector query failed: %s", e) return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py index b32bd13e8..f89c81c9d 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py @@ -46,7 +46,7 @@ def node_init(self): ) return super().node_init() except Exception as e: - log.error(f"Failed to initialize AnswerSynthesizeNode: {e}") + log.error("Failed to initialize AnswerSynthesizeNode: %s", e) from PyCGraph import CStatus return CStatus(-1, f"AnswerSynthesizeNode initialization failed: {e}") @@ -75,7 +75,7 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: if result.get("graph_vector_answer"): answer_types.append("graph_vector") - log.info(f"Answer synthesis completed for types: {', '.join(answer_types)}") + log.info("Answer synthesis completed for types: %s", ', '.join(answer_types)) # Print enabled answer types according to self.wk_input configuration wk_input_types = [] @@ -93,5 +93,5 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: return result except Exception as e: - log.error(f"Answer synthesis failed: {e}") + log.error("Answer synthesis failed: %s", e) return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py index ea77c6580..3f3facb68 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py @@ -47,7 +47,7 @@ def node_init(self): ) return super().node_init() except Exception as e: - log.error(f"Failed to initialize KeywordExtractNode: {e}") + log.error("Failed to initialize KeywordExtractNode: %s", e) return CStatus(-1, f"KeywordExtractNode initialization failed: {e}") def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: @@ -61,12 +61,12 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: log.warning("Keyword extraction result missing 'keywords' field") result["keywords"] = [] - log.info(f"Extracted keywords: {result.get('keywords', [])}") + log.info("Extracted keywords: %s", result.get('keywords', [])) return result except Exception as e: - log.error(f"Keyword extraction failed: {e}") + log.error("Keyword extraction failed: %s", e) # Add error flag to indicate failure error_result = data_json.copy() error_result["error"] = str(e) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index e49ca04ee..e3eea9f07 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -16,21 +16,16 @@ # under the License. -import asyncio import os from typing import Any, Dict, List, Optional import pandas as pd +from tqdm import tqdm -from hugegraph_llm.config import resource_path, huge_settings +from hugegraph_llm.config import resource_path from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.utils.embedding_utils import ( - get_embeddings_parallel, - get_filename_prefix, - get_index_folder_name, -) from hugegraph_llm.utils.log import log diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 51ff0b306..2bd4e2203 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -13,10 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import GParam, CStatus - from typing import Union, List, Optional, Any +from PyCGraph import GParam, CStatus + class WkFlowInput(GParam): texts: Union[str, List[str]] = None # texts input used by ChunkSplit Node From 3571450f458ffb88589a88ad72152cdde35595b5 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Sun, 12 Oct 2025 14:10:58 +0800 Subject: [PATCH 42/71] ci --- .../src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py | 2 +- .../src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py index f89c81c9d..cc3a7c25a 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py @@ -75,7 +75,7 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: if result.get("graph_vector_answer"): answer_types.append("graph_vector") - log.info("Answer synthesis completed for types: %s", ', '.join(answer_types)) + log.info("Answer synthesis completed for types: %s", ", ".join(answer_types)) # Print enabled answer types according to self.wk_input configuration wk_input_types = [] diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py index 3f3facb68..cb390b7b7 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py @@ -61,7 +61,7 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: log.warning("Keyword extraction result missing 'keywords' field") result["keywords"] = [] - log.info("Extracted keywords: %s", result.get('keywords', [])) + log.info("Extracted keywords: %s", result.get("keywords", [])) return result From 89151c2a607e05745b3ef72d6bb1c68efd7fdf77 Mon Sep 17 00:00:00 2001 From: jinglinwei Date: Mon, 13 Oct 2025 01:01:47 +0800 Subject: [PATCH 43/71] feat: support batch build gremlin examples & delete some doc related to Pipeline(old design) & refactor some operator's design and implementation --- README.md | 36 -- hugegraph-llm/README.md | 78 --- hugegraph-llm/pyproject.toml | 2 +- .../src/hugegraph_llm/api/rag_api.py | 57 ++- .../hugegraph_llm/demo/rag_demo/rag_block.py | 12 +- .../demo/rag_demo/text2gremlin_block.py | 155 ++---- .../flows/build_example_index.py | 59 +++ .../flows/get_graph_index_info.py | 6 +- .../src/hugegraph_llm/flows/graph_extract.py | 19 +- .../flows/rag_flow_graph_only.py | 110 +++-- .../flows/rag_flow_graph_vector.py | 23 +- .../src/hugegraph_llm/flows/scheduler.py | 29 +- .../flows/update_vid_embeddings.py | 2 +- .../src/hugegraph_llm/flows/utils.py | 2 +- .../src/hugegraph_llm/nodes/base_node.py | 9 +- .../nodes/hugegraph_node/fetch_graph_data.py | 6 +- .../nodes/hugegraph_node/graph_query_node.py | 459 ++++++++++++++++-- .../nodes/hugegraph_node/schema.py | 14 +- .../index_node/build_gremlin_example_index.py | 43 ++ .../index_node/semantic_id_query_node.py | 20 +- .../nodes/llm_node/text2gremlin.py | 18 +- hugegraph-llm/src/hugegraph_llm/nodes/util.py | 16 +- .../operators/gremlin_generate_task.py | 81 ---- .../hugegraph_op/commit_to_hugegraph.py | 59 ++- .../operators/hugegraph_op/graph_rag_query.py | 455 ----------------- .../index_op/build_semantic_index.py | 24 +- .../operators/kg_construction_task.py | 120 ----- .../{graph_rag_task.py => operator_list.py} | 222 +++++---- .../src/hugegraph_llm/operators/util.py | 27 -- .../src/hugegraph_llm/state/ai_state.py | 120 ++--- .../hugegraph_llm/utils/graph_index_utils.py | 80 +-- hugegraph-ml/pyproject.toml | 2 +- hugegraph-python-client/pyproject.toml | 2 +- pyproject.toml | 2 +- scripts/build_llm_image.sh | 2 +- .../hugegraph-llm/fixed_flow/design.md | 2 +- .../hugegraph-llm/fixed_flow/requirements.md | 0 .../hugegraph-llm/fixed_flow/tasks.md | 0 vermeer-python-client/pyproject.toml | 4 +- 39 files changed, 1064 insertions(+), 1313 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/build_example_index.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py delete mode 100644 hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py delete mode 100644 hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py delete mode 100644 hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py rename hugegraph-llm/src/hugegraph_llm/operators/{graph_rag_task.py => operator_list.py} (58%) delete mode 100644 hugegraph-llm/src/hugegraph_llm/operators/util.py mode change 100644 => 100755 scripts/build_llm_image.sh rename {.vibedev/spec => spec}/hugegraph-llm/fixed_flow/design.md (99%) rename {.vibedev/spec => spec}/hugegraph-llm/fixed_flow/requirements.md (100%) rename {.vibedev/spec => spec}/hugegraph-llm/fixed_flow/tasks.md (100%) diff --git a/README.md b/README.md index 14f02ca1c..a495968ec 100644 --- a/README.md +++ b/README.md @@ -75,42 +75,6 @@ python -m hugegraph_llm.demo.rag_demo.app > [!NOTE] > Examples assume you've activated the virtual environment with `source .venv/bin/activate` -#### GraphRAG - Question Answering - -```python -from hugegraph_llm.operators.graph_rag_task import RAGPipeline - -# Initialize RAG pipeline -graph_rag = RAGPipeline() - -# Ask questions about your graph -result = (graph_rag - .extract_keywords(text="Tell me about Al Pacino.") - .keywords_to_vid() - .query_graphdb(max_deep=2, max_graph_items=30) - .merge_dedup_rerank() - .synthesize_answer() - .run()) -``` - -#### Knowledge Graph Construction - -```python -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.kg_construction_task import KgBuilder - -# Build KG from text -TEXT = "Your text content here..." -builder = KgBuilder(LLMs().get_chat_llm()) - -(builder - .import_schema(from_hugegraph="hugegraph") - .chunk_split(TEXT) - .extract_info(extract_type="property_graph") - .commit_to_hugegraph() - .run()) -``` - #### Graph Machine Learning ```bash diff --git a/hugegraph-llm/README.md b/hugegraph-llm/README.md index 526320d4a..f0eeb3136 100644 --- a/hugegraph-llm/README.md +++ b/hugegraph-llm/README.md @@ -146,84 +146,6 @@ Use the Gradio interface for visual knowledge graph building: ![Knowledge Graph Builder](https://hugegraph.apache.org/docs/images/gradio-kg.png) -#### Programmatic Construction - -Build knowledge graphs with code using the `KgBuilder` class: - -```python -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.kg_construction_task import KgBuilder - -# Initialize and chain operations -TEXT = "Your input text here..." -builder = KgBuilder(LLMs().get_chat_llm()) - -( - builder - .import_schema(from_hugegraph="talent_graph").print_result() - .chunk_split(TEXT).print_result() - .extract_info(extract_type="property_graph").print_result() - .commit_to_hugegraph() - .run() -) -``` - -**Pipeline Workflow:** - -```mermaid -graph LR - A[Import Schema] --> B[Chunk Split] - B --> C[Extract Info] - C --> D[Commit to HugeGraph] - D --> E[Execute Pipeline] - - style A fill:#fff2cc - style B fill:#d5e8d4 - style C fill:#dae8fc - style D fill:#f8cecc - style E fill:#e1d5e7 -``` - -### Graph-Enhanced RAG - -Leverage HugeGraph for retrieval-augmented generation: - -```python -from hugegraph_llm.operators.graph_rag_task import RAGPipeline - -# Initialize RAG pipeline -graph_rag = RAGPipeline() - -# Execute RAG workflow -( - graph_rag - .extract_keywords(text="Tell me about Al Pacino.") - .keywords_to_vid() - .query_graphdb(max_deep=2, max_graph_items=30) - .merge_dedup_rerank() - .synthesize_answer(vector_only_answer=False, graph_only_answer=True) - .run(verbose=True) -) -``` - -**RAG Pipeline Flow:** - -```mermaid -graph TD - A[User Query] --> B[Extract Keywords] - B --> C[Match Graph Nodes] - C --> D[Retrieve Graph Context] - D --> E[Rerank Results] - E --> F[Generate Answer] - - style A fill:#e3f2fd - style B fill:#f3e5f5 - style C fill:#e8f5e8 - style D fill:#fff3e0 - style E fill:#fce4ec - style F fill:#e0f2f1 -``` - ## 🔧 Configuration After running the demo, configuration files are automatically generated: diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index 09c49ae26..f5301591b 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -17,7 +17,7 @@ [project] name = "hugegraph-llm" -version = "1.5.0" +version = "1.7.0" description = "A tool for the implementation and research related to large language models." authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py index 5c9295efa..707069745 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py @@ -31,8 +31,8 @@ from hugegraph_llm.config import huge_settings from hugegraph_llm.api.models.rag_response import RAGResponse from hugegraph_llm.config import llm_settings, prompt +from hugegraph_llm.utils.graph_index_utils import get_vertex_details from hugegraph_llm.utils.log import log -from hugegraph_llm.flows.scheduler import SchedulerSingleton # pylint: disable=too-many-statements @@ -67,7 +67,8 @@ def rag_answer_api(req: RAGRequest): # Keep prompt params in the end custom_related_information=req.custom_priority_info, answer_prompt=req.answer_prompt or prompt.answer_prompt, - keywords_extract_prompt=req.keywords_extract_prompt or prompt.keywords_extract_prompt, + keywords_extract_prompt=req.keywords_extract_prompt + or prompt.keywords_extract_prompt, gremlin_prompt=req.gremlin_prompt or prompt.gremlin_generate_prompt, ) # TODO: we need more info in the response for users to understand the query logic @@ -76,7 +77,8 @@ def rag_answer_api(req: RAGRequest): **{ key: value for key, value in zip( - ["raw_answer", "vector_only", "graph_only", "graph_vector_answer"], result + ["raw_answer", "vector_only", "graph_only", "graph_vector_answer"], + result, ) if getattr(req, key) }, @@ -110,12 +112,7 @@ def graph_rag_recall_api(req: GraphRAGRequest): ) if req.get_vertex_only: - from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery - - graph_rag = GraphRAGQuery() - graph_rag.init_client(result) - vertex_details = graph_rag.get_vertex_details(result["match_vids"]) - + vertex_details = get_vertex_details(result["match_vids"], result) if vertex_details: result["match_vids"] = vertex_details @@ -135,7 +132,9 @@ def graph_rag_recall_api(req: GraphRAGRequest): except TypeError as e: log.error("TypeError in graph_rag_recall_api: %s", e) - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) + ) from e except Exception as e: log.error("Unexpected error occurred: %s", e) raise HTTPException( @@ -146,7 +145,9 @@ def graph_rag_recall_api(req: GraphRAGRequest): @router.post("/config/graph", status_code=status.HTTP_201_CREATED) def graph_config_api(req: GraphConfigRequest): # Accept status code - res = apply_graph_conf(req.url, req.name, req.user, req.pwd, req.gs, origin_call="http") + res = apply_graph_conf( + req.url, req.name, req.user, req.pwd, req.gs, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) # TODO: restructure the implement of llm to three types, like "/config/chat_llm" @@ -156,10 +157,16 @@ def llm_config_api(req: LLMConfigRequest): if req.llm_type == "openai": res = apply_llm_conf( - req.api_key, req.api_base, req.language_model, req.max_tokens, origin_call="http" + req.api_key, + req.api_base, + req.language_model, + req.max_tokens, + origin_call="http", ) else: - res = apply_llm_conf(req.host, req.port, req.language_model, None, origin_call="http") + res = apply_llm_conf( + req.host, req.port, req.language_model, None, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/embedding", status_code=status.HTTP_201_CREATED) @@ -171,7 +178,9 @@ def embedding_config_api(req: LLMConfigRequest): req.api_key, req.api_base, req.language_model, origin_call="http" ) else: - res = apply_embedding_conf(req.host, req.port, req.language_model, origin_call="http") + res = apply_embedding_conf( + req.host, req.port, req.language_model, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/rerank", status_code=status.HTTP_201_CREATED) @@ -183,7 +192,9 @@ def rerank_config_api(req: RerankerConfigRequest): req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http" ) elif req.reranker_type == "siliconflow": - res = apply_reranker_conf(req.api_key, req.reranker_model, None, origin_call="http") + res = apply_reranker_conf( + req.api_key, req.reranker_model, None, origin_call="http" + ) else: res = status.HTTP_501_NOT_IMPLEMENTED return generate_response(RAGResponse(status_code=res, message="Missing Value")) @@ -196,20 +207,20 @@ def text2gremlin_api(req: GremlinGenerateRequest): # Basic parameter validation: empty query => 400 if not req.query or not str(req.query).strip(): raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, detail="Query must not be empty." + status_code=status.HTTP_400_BAD_REQUEST, + detail="Query must not be empty.", ) output_types_str_list = None if req.output_types: output_types_str_list = [ot.value for ot in req.output_types] - response_dict = SchedulerSingleton.get_instance().schedule_flow( - "text2gremlin", - req.query, - req.example_num, - huge_settings.graph_name, - req.gremlin_prompt, - output_types_str_list, + response_dict = gremlin_generate_selective_func( + inp=req.query, + example_num=req.example_num, + schema_input=huge_settings.graph_name, + gremlin_prompt_input=req.gremlin_prompt, + requested_outputs=output_types_str_list, ) return response_dict except HTTPException as e: diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 5ff3df931..60ca6ae55 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -51,11 +51,7 @@ def rag_answer( ) -> Tuple: """ Generate an answer using the RAG (Retrieval-Augmented Generation) pipeline. - 1. Initialize the RAGPipeline. - 2. Select vector search or graph search based on parameters. - 3. Merge, deduplicate, and rerank the results. - 4. Synthesize the final answer. - 5. Run the pipeline and return the results. + Fetch the Scheduler to deal with the request """ graph_search, gremlin_prompt, vector_search = update_ui_configs( answer_prompt, @@ -172,11 +168,7 @@ async def rag_answer_streaming( ) -> AsyncGenerator[Tuple[str, str, str, str], None]: """ Generate an answer using the RAG (Retrieval-Augmented Generation) pipeline. - 1. Initialize the RAGPipeline. - 2. Select vector search or graph search based on parameters. - 3. Merge, deduplicate, and rerank the results. - 4. Synthesize the final answer. - 5. Run the pipeline and return the results. + Fetch the Scheduler to deal with the request """ graph_search, gremlin_prompt, vector_search = update_ui_configs( answer_prompt, diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 6600d7c41..04aef1c77 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -25,11 +25,6 @@ import pandas as pd from hugegraph_llm.config import prompt, resource_path, huge_settings -from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.graph_rag_task import RAGPipeline -from hugegraph_llm.operators.gremlin_generate_task import GremlinGenerator -from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager from hugegraph_llm.utils.embedding_utils import get_index_folder_name from hugegraph_llm.utils.hugegraph_utils import run_gremlin_query from hugegraph_llm.utils.log import log @@ -86,7 +81,9 @@ def store_schema(schema, question, gremlin_prompt): def build_example_vector_index(temp_file) -> dict: - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) index_path = os.path.join(resource_path, folder_name, "gremlin_examples") if not os.path.exists(index_path): os.makedirs(index_path) @@ -98,7 +95,9 @@ def build_example_vector_index(temp_file) -> dict: timestamp = datetime.now().strftime("%Y%m%d%H%M%S") _, file_name = os.path.split(f"{name}_{timestamp}{ext}") log.info("Copying file to: %s", file_name) - target_file = os.path.join(resource_path, folder_name, "gremlin_examples", file_name) + target_file = os.path.join( + resource_path, folder_name, "gremlin_examples", file_name + ) try: import shutil @@ -116,11 +115,10 @@ def build_example_vector_index(temp_file) -> dict: else: log.critical("Unsupported file format. Please input a JSON or CSV file.") return {"error": "Unsupported file format. Please input a JSON or CSV file."} - builder = GremlinGenerator( - llm=LLMs().get_text2gql_llm(), - embedding=Embeddings().get_embedding(), + + return SchedulerSingleton.get_instance().schedule_flow( + "build_examples_index", examples ) - return builder.example_index_build(examples).run() def _process_schema(schema, generator, sm): @@ -182,43 +180,6 @@ def _execute_queries(context, output_types): context["raw_exec_res"] = "" -def gremlin_generate( - inp, example_num, schema, gremlin_prompt, requested_outputs: Optional[List[str]] = None -) -> GremlinResult: - generator = GremlinGenerator( - llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() - ) - sm = SchemaManager(graph_name=schema) - - processed_schema, short_schema = _process_schema(schema, generator, sm) - if processed_schema is None and short_schema is None: - return GremlinResult.error("Invalid JSON schema, please check the format carefully.") - - updated_schema = sm.simple_schema(processed_schema) if short_schema else processed_schema - store_schema(str(updated_schema), inp, gremlin_prompt) - - output_types = _configure_output_types(requested_outputs) - - context = ( - generator.example_index_query(example_num) - .gremlin_generate_synthesize(updated_schema, gremlin_prompt) - .run(query=inp) - ) - - _execute_queries(context, output_types) - - match_result = json.dumps( - context.get("match_result", "No Results"), ensure_ascii=False, indent=2 - ) - return GremlinResult.success_result( - match_result=match_result, - template_gremlin=context["result"], - raw_gremlin=context["raw_result"], - template_exec=context["template_exec_res"], - raw_exec=context["raw_exec_res"], - ) - - def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: mini_schema = {} @@ -226,7 +187,11 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} + new_vertex = { + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex + } mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) @@ -305,15 +270,21 @@ def create_text2gremlin_block() -> Tuple: with gr.Row(): with gr.Column(scale=1): input_box = gr.Textbox( - value=prompt.default_question, label="Nature Language Query", show_copy_button=True + value=prompt.default_question, + label="Nature Language Query", + show_copy_button=True, ) match = gr.Code( label="Similar Template (TopN)", language="javascript", elem_classes="code-container-show", ) - initialized_out = gr.Textbox(label="Gremlin With Template", show_copy_button=True) - raw_out = gr.Textbox(label="Gremlin Without Template", show_copy_button=True) + initialized_out = gr.Textbox( + label="Gremlin With Template", show_copy_button=True + ) + raw_out = gr.Textbox( + label="Gremlin Without Template", show_copy_button=True + ) tmpl_exec_out = gr.Code( label="Query With Template Output", language="json", @@ -330,7 +301,10 @@ def create_text2gremlin_block() -> Tuple: minimum=0, maximum=10, step=1, value=2, label="Number of refer examples" ) schema_box = gr.Textbox( - value=prompt.text2gql_graph_schema, label="Schema", lines=2, show_copy_button=True + value=prompt.text2gql_graph_schema, + label="Schema", + lines=2, + show_copy_button=True, ) prompt_box = gr.Textbox( value=prompt.gremlin_generate_prompt, @@ -362,24 +336,21 @@ def graph_rag_recall( get_vertex_only: bool = False, ) -> dict: store_schema(prompt.text2gql_graph_schema, query, gremlin_prompt) - rag = RAGPipeline() - rag.extract_keywords().keywords_to_vid( + context = SchedulerSingleton.get_instance().schedule_flow( + "rag_graph_only", + query=query, + gremlin_tmpl_num=gremlin_tmpl_num, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + gremlin_prompt=gremlin_prompt, + max_graph_items=max_graph_items, + topk_return_results=topk_return_results, vector_dis_threshold=vector_dis_threshold, topk_per_keyword=topk_per_keyword, + is_graph_rag_recall=True, + is_vector_only=get_vertex_only, ) - - if not get_vertex_only: - rag.import_schema(huge_settings.graph_name).query_graphdb( - num_gremlin_generate_example=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - max_graph_items=max_graph_items, - ).merge_dedup_rerank( - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - custom_related_information=custom_related_information, - topk_return_results=topk_return_results, - ) - context = rag.run(verbose=True, query=query, graph_search=True) return context @@ -390,45 +361,13 @@ def gremlin_generate_selective( gremlin_prompt_input: str, requested_outputs: Optional[List[str]] = None, ) -> Dict[str, Any]: - """ - Wraps the gremlin_generate function to return a dictionary of outputs - based on the requested_outputs list of strings. - """ - output_keys = [ - "match_result", - "template_gremlin", - "raw_gremlin", - "template_execution_result", - "raw_execution_result", - ] - if not requested_outputs: # None or empty list - requested_outputs = output_keys - - result = gremlin_generate( - inp, example_num, schema_input, gremlin_prompt_input, requested_outputs + response_dict = SchedulerSingleton.get_instance().schedule_flow( + "text2gremlin", + inp, + example_num, + schema_input, + gremlin_prompt_input, + requested_outputs, ) - outputs_dict: Dict[str, Any] = {} - - if not result.success: - # Handle error case - if "match_result" in requested_outputs: - outputs_dict["match_result"] = result.match_result - if result.error_message: - outputs_dict["error_detail"] = result.error_message - return outputs_dict - - # Handle successful case - output_mapping = { - "match_result": result.match_result, - "template_gremlin": result.template_gremlin, - "raw_gremlin": result.raw_gremlin, - "template_execution_result": result.template_exec_result, - "raw_execution_result": result.raw_exec_result, - } - - for key in requested_outputs: - if key in output_mapping: - outputs_dict[key] = output_mapping[key] - - return outputs_dict + return response_dict diff --git a/hugegraph-llm/src/hugegraph_llm/flows/build_example_index.py b/hugegraph-llm/src/hugegraph_llm/flows/build_example_index.py new file mode 100644 index 000000000..903059fbf --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/build_example_index.py @@ -0,0 +1,59 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Dict, Optional + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.nodes.index_node.build_gremlin_example_index import ( + BuildGremlinExampleIndexNode, +) +from hugegraph_llm.utils.log import log + +import json +from PyCGraph import GPipeline + + +class BuildExampleIndexFlow(BaseFlow): + def __init__(self): + pass + + def prepare( + self, prepared_input: WkFlowInput, examples: Optional[List[Dict[str, str]]] + ): + prepared_input.examples = examples + return + + def build_flow(self, examples=None): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, examples=examples) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + build_node = BuildGremlinExampleIndexNode() + pipeline.registerGElement(build_node, set(), "build_examples_index") + + return pipeline + + def post_deal(self, pipeline=None): + state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + try: + formatted_schema = json.dumps(state_json, ensure_ascii=False, indent=2) + return formatted_schema + except (TypeError, ValueError) as e: + log.error("Failed to format schema: %s", e) + return str(state_json) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py index 7d2735352..439c3a346 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py @@ -48,7 +48,9 @@ def build_flow(self, *args, **kwargs): def post_deal(self, pipeline=None): graph_summary_info = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) filename_prefix = get_filename_prefix( llm_settings.embedding_type, @@ -56,7 +58,7 @@ def post_deal(self, pipeline=None): ) try: vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) - except FileNotFoundError: + except (RuntimeError, OSError): return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) graph_summary_info["vid_index"] = { "embed_dim": vector_index.index.d, diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py index 55f53b7ad..89a5fa5f8 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py @@ -27,21 +27,30 @@ class GraphExtractFlow(BaseFlow): def __init__(self): pass - def prepare(self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type): + def prepare( + self, + prepared_input: WkFlowInput, + schema, + texts, + example_prompt, + extract_type, + language="zh", + ): # prepare input data prepared_input.texts = texts - prepared_input.language = "zh" + prepared_input.language = language prepared_input.split_type = "document" prepared_input.example_prompt = example_prompt prepared_input.schema = schema prepared_input.extract_type = extract_type - return - def build_flow(self, schema, texts, example_prompt, extract_type): + def build_flow(self, schema, texts, example_prompt, extract_type, language="zh"): pipeline = GPipeline() prepared_input = WkFlowInput() # prepare input data - self.prepare(prepared_input, schema, texts, example_prompt, extract_type) + self.prepare( + prepared_input, schema, texts, example_prompt, extract_type, language + ) pipeline.createGParam(prepared_input, "wkflow_input") pipeline.createGParam(WkFlowState(), "wkflow_state") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py index 5feb3d471..c4cfd46dc 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py @@ -13,11 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json -from typing import Optional, Literal +from typing import Optional, Literal, cast -from PyCGraph import GPipeline +from PyCGraph import GPipeline, GRegion, GCondition from hugegraph_llm.flows.common import BaseFlow from hugegraph_llm.nodes.llm_node.keyword_extract_node import KeywordExtractNode @@ -31,6 +30,22 @@ from hugegraph_llm.utils.log import log +class GraphRecallCondition(GCondition): + def choose(self): + prepare_input: WkFlowInput = cast( + WkFlowInput, self.getGParamWithNoEmpty("wkflow_input") + ) + return 0 if prepare_input.is_graph_rag_recall else 1 + + +class VectorOnlyCondition(GCondition): + def choose(self): + prepare_input: WkFlowInput = cast( + WkFlowInput, self.getGParamWithNoEmpty("wkflow_input") + ) + return 0 if prepare_input.is_vector_only else 1 + + class RAGGraphOnlyFlow(BaseFlow): """ Workflow for graph-only answering (graph_only_answer) @@ -40,13 +55,12 @@ def prepare( self, prepared_input: WkFlowInput, query: str, - vector_search: bool = None, - graph_search: bool = None, - raw_answer: bool = None, - vector_only_answer: bool = None, - graph_only_answer: bool = None, - graph_vector_answer: bool = None, - graph_ratio: float = 0.5, + vector_search: bool = False, + graph_search: bool = True, + raw_answer: bool = False, + vector_only_answer: bool = False, + graph_only_answer: bool = True, + graph_vector_answer: bool = False, rerank_method: Literal["bleu", "reranker"] = "bleu", near_neighbor_first: bool = False, custom_related_information: str = "", @@ -54,10 +68,12 @@ def prepare( keywords_extract_prompt: Optional[str] = None, gremlin_tmpl_num: Optional[int] = -1, gremlin_prompt: Optional[str] = None, - max_graph_items: int = None, - topk_return_results: int = None, - vector_dis_threshold: float = None, - topk_per_keyword: int = None, + max_graph_items: Optional[int] = None, + topk_return_results: Optional[int] = None, + vector_dis_threshold: Optional[float] = None, + topk_per_keyword: Optional[int] = None, + is_graph_rag_recall: bool = False, + is_vector_only: bool = False, **_: dict, ): prepared_input.query = query @@ -90,6 +106,8 @@ def prepare( ) prepared_input.schema = huge_settings.graph_name + prepared_input.is_graph_rag_recall = is_graph_rag_recall + prepared_input.is_vector_only = is_vector_only prepared_input.data_json = { "query": query, "vector_search": vector_search, @@ -106,48 +124,56 @@ def build_flow(self, **kwargs): pipeline.createGParam(WkFlowState(), "wkflow_state") # Create nodes and register them with registerGElement - only_keyword_extract_node = KeywordExtractNode() - only_semantic_id_query_node = SemanticIdQueryNode() + only_keyword_extract_node = KeywordExtractNode("only_keyword") + only_semantic_id_query_node = SemanticIdQueryNode( + {only_keyword_extract_node}, "only_semantic" + ) + vector_region: GRegion = GRegion( + [only_keyword_extract_node, only_semantic_id_query_node] + ) + only_schema_node = SchemaNode() - only_graph_query_node = GraphQueryNode() - merge_rerank_node = MergeRerankNode() + schema_node = VectorOnlyCondition([GRegion(), only_schema_node]) + only_graph_query_node = GraphQueryNode("only_graph") + merge_rerank_node = MergeRerankNode({only_graph_query_node}, "merge_rerank") + graph_region: GRegion = GRegion([only_graph_query_node, merge_rerank_node]) + graph_condition_region = VectorOnlyCondition([GRegion(), graph_region]) + answer_synthesize_node = AnswerSynthesizeNode() + answer_node = GraphRecallCondition([GRegion(), answer_synthesize_node]) - pipeline.registerGElement(only_keyword_extract_node, set(), "only_keyword") - pipeline.registerGElement( - only_semantic_id_query_node, {only_keyword_extract_node}, "only_semantic" - ) - pipeline.registerGElement(only_schema_node, set(), "only_schema") + pipeline.registerGElement(vector_region, set(), "vector_fetch") + pipeline.registerGElement(schema_node, set(), "schema_condition") pipeline.registerGElement( - only_graph_query_node, - {only_schema_node, only_semantic_id_query_node}, - "only_graph", + graph_condition_region, + {schema_node, vector_region}, + "graph_condition", ) pipeline.registerGElement( - merge_rerank_node, {only_graph_query_node}, "merge_one" + answer_node, {graph_condition_region}, "answer_condition" ) - pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "graph") log.info("RAGGraphOnlyFlow pipeline built successfully") return pipeline def post_deal(self, pipeline=None): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) + return {"error": "No pipeline provided"} try: + prepare_input = cast( + WkFlowInput, pipeline.getGParamWithNoEmpty("wkflow_input") + ) res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() log.info("RAGGraphOnlyFlow post processing success") - return { - "raw_answer": res.get("raw_answer", ""), - "vector_only_answer": res.get("vector_only_answer", ""), - "graph_only_answer": res.get("graph_only_answer", ""), - "graph_vector_answer": res.get("graph_vector_answer", ""), - } + return ( + { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } + if not prepare_input.is_graph_rag_recall + else res + ) except Exception as e: log.error(f"RAGGraphOnlyFlow post processing failed: {e}") - return json.dumps( - {"error": f"Post processing failed: {str(e)}"}, - ensure_ascii=False, - indent=2, - ) + return {"error": f"Post processing failed: {str(e)}"} diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py index 2f4a2bfa2..9fd4d96e0 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json from typing import Optional, Literal @@ -41,12 +40,12 @@ def prepare( self, prepared_input: WkFlowInput, query: str, - vector_search: bool = None, - graph_search: bool = None, - raw_answer: bool = None, - vector_only_answer: bool = None, - graph_only_answer: bool = None, - graph_vector_answer: bool = None, + vector_search: bool = True, + graph_search: bool = True, + raw_answer: bool = False, + vector_only_answer: bool = False, + graph_only_answer: bool = False, + graph_vector_answer: bool = True, graph_ratio: float = 0.5, rerank_method: Literal["bleu", "reranker"] = "bleu", near_neighbor_first: bool = False, @@ -137,9 +136,7 @@ def build_flow(self, **kwargs): def post_deal(self, pipeline=None): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) + return {"error": "No pipeline provided"} try: res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() log.info("RAGGraphVectorFlow post processing success") @@ -151,8 +148,4 @@ def post_deal(self, pipeline=None): } except Exception as e: log.error(f"RAGGraphVectorFlow post processing failed: {e}") - return json.dumps( - {"error": f"Post processing failed: {str(e)}"}, - ensure_ascii=False, - indent=2, - ) + return {"error": f"Post processing failed: {str(e)}"} diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index 5afa1bf8e..7cc5653d5 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -18,9 +18,10 @@ from PyCGraph import GPipeline, GPipelineManager from hugegraph_llm.flows.build_vector_index import BuildVectorIndexFlow from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.flows.build_example_index import BuildExampleIndexFlow from hugegraph_llm.flows.graph_extract import GraphExtractFlow from hugegraph_llm.flows.import_graph_data import ImportGraphDataFlow -from hugegraph_llm.flows.update_vid_embeddings import UpdateVidEmbeddingsFlows +from hugegraph_llm.flows.update_vid_embeddings import UpdateVidEmbeddingsFlow from hugegraph_llm.flows.get_graph_index_info import GetGraphIndexInfoFlow from hugegraph_llm.flows.build_schema import BuildSchemaFlow from hugegraph_llm.flows.prompt_generate import PromptGenerateFlow @@ -54,7 +55,7 @@ def __init__(self, max_pipeline: int = 10): } self.pipeline_pool["update_vid_embeddings"] = { "manager": GPipelineManager(), - "flow": UpdateVidEmbeddingsFlows(), + "flow": UpdateVidEmbeddingsFlow(), } self.pipeline_pool["get_graph_index_info"] = { "manager": GPipelineManager(), @@ -89,17 +90,21 @@ def __init__(self, max_pipeline: int = 10): "manager": GPipelineManager(), "flow": RAGGraphVectorFlow(), } + self.pipeline_pool["build_examples_index"] = { + "manager": GPipelineManager(), + "flow": BuildExampleIndexFlow(), + } self.max_pipeline = max_pipeline # TODO: Implement Agentic Workflow def agentic_flow(self): pass - def schedule_flow(self, flow: str, *args, **kwargs): - if flow not in self.pipeline_pool: - raise ValueError(f"Unsupported workflow {flow}") - manager: GPipelineManager = self.pipeline_pool[flow]["manager"] - flow: BaseFlow = self.pipeline_pool[flow]["flow"] + def schedule_flow(self, flow_name: str, *args, **kwargs): + if flow_name not in self.pipeline_pool: + raise ValueError(f"Unsupported workflow {flow_name}") + manager: GPipelineManager = self.pipeline_pool[flow_name]["manager"] + flow: BaseFlow = self.pipeline_pool[flow_name]["flow"] pipeline: GPipeline = manager.fetch() if pipeline is None: # call coresponding flow_func to create new workflow @@ -130,11 +135,11 @@ def schedule_flow(self, flow: str, *args, **kwargs): manager.release(pipeline) return res - async def schedule_stream_flow(self, flow: str, *args, **kwargs): - if flow not in self.pipeline_pool: - raise ValueError(f"Unsupported workflow {flow}") - manager: GPipelineManager = self.pipeline_pool[flow]["manager"] - flow: BaseFlow = self.pipeline_pool[flow]["flow"] + async def schedule_stream_flow(self, flow_name: str, *args, **kwargs): + if flow_name not in self.pipeline_pool: + raise ValueError(f"Unsupported workflow {flow_name}") + manager: GPipelineManager = self.pipeline_pool[flow_name]["manager"] + flow: BaseFlow = self.pipeline_pool[flow_name]["flow"] pipeline: GPipeline = manager.fetch() if pipeline is None: # call coresponding flow_func to create new workflow 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 b3f0d9923..01d14a0e2 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py @@ -20,7 +20,7 @@ from hugegraph_llm.state.ai_state import WkFlowState -class UpdateVidEmbeddingsFlows(BaseFlow): +class UpdateVidEmbeddingsFlow(BaseFlow): def prepare(self, prepared_input: WkFlowInput): return CStatus() diff --git a/hugegraph-llm/src/hugegraph_llm/flows/utils.py b/hugegraph-llm/src/hugegraph_llm/flows/utils.py index b4ba05c84..5bd7dc73c 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/utils.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/utils.py @@ -19,7 +19,7 @@ from hugegraph_llm.utils.log import log -def prepare_schema(prepared_input: WkFlowInput, schema): +def prepare_schema(prepared_input: WkFlowInput, schema: str) -> None: schema = schema.strip() if schema.startswith("{"): try: diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py index f90167305..9c6457b51 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py @@ -16,6 +16,7 @@ from PyCGraph import GNode, CStatus from hugegraph_llm.nodes.util import init_context from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.utils.log import log class BaseNode(GNode): @@ -43,6 +44,8 @@ def run(self): sts = self.node_init() if sts.isErr(): return sts + if self.context is None: + return CStatus(-1, "Context not initialized") self.context.lock() try: data_json = self.context.to_json() @@ -60,8 +63,10 @@ def run(self): self.context.lock() try: - if isinstance(res, dict): + if res is not None and isinstance(res, dict): self.context.assign_from_json(res) + elif res is not None: + log.warning(f"operator_schedule returned non-dict type: {type(res)}") finally: self.context.unlock() return CStatus() @@ -69,6 +74,6 @@ def run(self): def operator_schedule(self, data_json): """ Interface for scheduling the operator, can be overridden by subclasses. - Returns a CStatus object indicating whether scheduling succeeded. + Subclasses should return a dict to update the workflow state, or None to skip state update. """ pass 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 99b428e5e..078d7fdab 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 @@ -25,7 +25,11 @@ class FetchGraphDataNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): - self.fetch_graph_data_op = FetchGraphData(get_hg_client()) + try: + client = get_hg_client() + except Exception as e: + raise RuntimeError(f"can't initalzie HugeGraph client: {e}") + self.fetch_graph_data_op = FetchGraphData(client) return super().node_init() def operator_schedule(self, data_json): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py index ae65ccb33..14f773db5 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py @@ -13,12 +13,63 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json + from PyCGraph import CStatus -from typing import Dict, Any +from typing import Dict, Any, Tuple, List, Set, Optional from hugegraph_llm.nodes.base_node import BaseNode -from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.operators.operator_list import OperatorList from hugegraph_llm.utils.log import log +from pyhugegraph.client import PyHugeClient + +# TODO: remove 'as('subj)' step +VERTEX_QUERY_TPL = "g.V({keywords}).limit(8).as('subj').toList()" + +# TODO: we could use a simpler query (like kneighbor-api to get the edges) +# TODO: test with profile()/explain() to speed up the query +VID_QUERY_NEIGHBOR_TPL = """\ +g.V({keywords}) +.repeat( + bothE({edge_labels}).limit({edge_limit}).otherV().dedup() +).times({max_deep}).emit() +.simplePath() +.path() +.by(project('label', 'id', 'props') + .by(label()) + .by(id()) + .by(valueMap().by(unfold())) +) +.by(project('label', 'inV', 'outV', 'props') + .by(label()) + .by(inV().id()) + .by(outV().id()) + .by(valueMap().by(unfold())) +) +.limit({max_items}) +.toList() +""" + +PROPERTY_QUERY_NEIGHBOR_TPL = """\ +g.V().has('{prop}', within({keywords})) +.repeat( + bothE({edge_labels}).limit({edge_limit}).otherV().dedup() +).times({max_deep}).emit() +.simplePath() +.path() +.by(project('label', 'props') + .by(label()) + .by(valueMap().by(unfold())) +) +.by(project('label', 'inV', 'outV', 'props') + .by(label()) + .by(inV().values('{prop}')) + .by(outV().values('{prop}')) + .by(valueMap().by(unfold())) +) +.limit({max_items}) +.toList() +""" class GraphQueryNode(BaseNode): @@ -26,39 +77,32 @@ class GraphQueryNode(BaseNode): Graph query node, responsible for retrieving relevant information from the graph database. """ - graph_rag_query: GraphRAGQuery - def node_init(self): """ Initialize the graph query operator. """ try: - graph_name = huge_settings.graph_name - if not graph_name: - return CStatus(-1, "graph_name is required in wk_input") - - max_deep = self.wk_input.max_deep or 2 - max_graph_items = ( + self._client: PyHugeClient = 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, + ) + self._max_deep = self.wk_input.max_deep or 2 + self._max_items = ( self.wk_input.max_graph_items or huge_settings.max_graph_items ) - max_v_prop_len = self.wk_input.max_v_prop_len or 2048 - max_e_prop_len = self.wk_input.max_e_prop_len or 256 - prop_to_match = self.wk_input.prop_to_match - num_gremlin_generate_example = self.wk_input.gremlin_tmpl_num or -1 - gremlin_prompt = ( + self._prop_to_match = self.wk_input.prop_to_match + self._num_gremlin_generate_example = self.wk_input.gremlin_tmpl_num or -1 + self.gremlin_prompt = ( self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt ) - - # Initialize GraphRAGQuery operator - self.graph_rag_query = GraphRAGQuery( - max_deep=max_deep, - max_graph_items=max_graph_items, - max_v_prop_len=max_v_prop_len, - max_e_prop_len=max_e_prop_len, - prop_to_match=prop_to_match, - num_gremlin_generate_example=num_gremlin_generate_example, - gremlin_prompt=gremlin_prompt, - ) + self._limit_property = huge_settings.limit_property.lower() == "true" + self._max_v_prop_len = self.wk_input.max_v_prop_len or 2048 + self._max_e_prop_len = self.wk_input.max_e_prop_len or 256 + self._schema = "" + self.operator_list = OperatorList(None, None) return super().node_init() except Exception as e: @@ -66,6 +110,350 @@ def node_init(self): return CStatus(-1, f"GraphQueryNode initialization failed: {e}") + # TODO: move this method to a util file for reuse (remove self param) + def init_client(self, context): + """Initialize the HugeGraph client from context or default settings.""" + # pylint: disable=R0915 (too-many-statements) + if self._client is None: + if isinstance(context.get("graph_client"), PyHugeClient): + self._client = context["graph_client"] + else: + url = context.get("url") or "http://localhost:8080" + graph = context.get("graph") or "hugegraph" + user = context.get("user") or "admin" + pwd = context.get("pwd") or "admin" + gs = context.get("graphspace") or None + self._client = PyHugeClient(url, graph, user, pwd, gs) + assert self._client is not None, "No valid graph to search." + + def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: + query = context["query"] + vertices = context.get("match_vids") + query_embedding = context.get("query_embedding") + + self.operator_list.clear() + self.operator_list.example_index_query( + num_examples=self._num_gremlin_generate_example + ) + gremlin_response = self.operator_list.gremlin_generate_synthesize( + context["simple_schema"], + vertices=vertices, + gremlin_prompt=self.gremlin_prompt, + ).run(query=query, query_embedding=query_embedding) + if self._num_gremlin_generate_example > 0: + gremlin = gremlin_response["result"] + else: + gremlin = gremlin_response["raw_result"] + log.info("Generated gremlin: %s", gremlin) + context["gremlin"] = gremlin + try: + result = self._client.gremlin().exec(gremlin=gremlin)["data"] + if result == [None]: + result = [] + context["graph_result"] = [ + json.dumps(item, ensure_ascii=False) for item in result + ] + if context["graph_result"]: + context["graph_result_flag"] = 1 + context["graph_context_head"] = ( + f"The following are graph query result " + f"from gremlin query `{gremlin}`.\n" + ) + except Exception as e: # pylint: disable=broad-except + log.error(e) + context["graph_result"] = "" + return context + + def _limit_property_query( + self, value: Optional[str], item_type: str + ) -> Optional[str]: + # NOTE: we skip the filter for list/set type (e.g., list of string, add it if needed) + if not self._limit_property or not isinstance(value, str): + return value + + max_len = self._max_v_prop_len if item_type == "v" else self._max_e_prop_len + return value[:max_len] if value else value + + def _process_vertex( + self, + item: Any, + flat_rel: str, + node_cache: Set[str], + prior_edge_str_len: int, + depth: int, + nodes_with_degree: List[str], + use_id_to_match: bool, + v_cache: Set[str], + ) -> Tuple[str, int, int]: + matched_str = ( + item["id"] if use_id_to_match else item["props"][self._prop_to_match] + ) + if matched_str in node_cache: + flat_rel = flat_rel[:-prior_edge_str_len] + return flat_rel, prior_edge_str_len, depth + + node_cache.add(matched_str) + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'v')}" + for k, v in item["props"].items() + if v + ) + + # TODO: we may remove label id or replace with label name + if matched_str in v_cache: + node_str = matched_str + else: + v_cache.add(matched_str) + node_str = f"{item['id']}{{{props_str}}}" + + flat_rel += node_str + nodes_with_degree.append(node_str) + depth += 1 + return flat_rel, prior_edge_str_len, depth + + def _process_edge( + self, + item: Any, + path_str: str, + raw_flat_rel: List[Any], + i: int, + use_id_to_match: bool, + e_cache: Set[Tuple[str, str, str]], + ) -> Tuple[str, int]: + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'e')}" + for k, v in item["props"].items() + if v + ) + props_str = f"{{{props_str}}}" if props_str else "" + prev_matched_str = ( + raw_flat_rel[i - 1]["id"] + if use_id_to_match + else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] + ) + + edge_key = (item["inV"], item["label"], item["outV"]) + if edge_key not in e_cache: + e_cache.add(edge_key) + edge_label = f"{item['label']}{props_str}" + else: + edge_label = item["label"] + + edge_str = ( + f"--[{edge_label}]-->" + if item["outV"] == prev_matched_str + else f"<--[{edge_label}]--" + ) + path_str += edge_str + prior_edge_str_len = len(edge_str) + return path_str, prior_edge_str_len + + def _process_path( + self, + path: Any, + use_id_to_match: bool, + v_cache: Set[str], + e_cache: Set[Tuple[str, str, str]], + ) -> Tuple[str, List[str]]: + flat_rel = "" + raw_flat_rel = path["objects"] + assert len(raw_flat_rel) % 2 == 1, "The length of raw_flat_rel should be odd." + + node_cache = set() + prior_edge_str_len = 0 + depth = 0 + nodes_with_degree = [] + + for i, item in enumerate(raw_flat_rel): + if i % 2 == 0: + # Process each vertex + flat_rel, prior_edge_str_len, depth = self._process_vertex( + item, + flat_rel, + node_cache, + prior_edge_str_len, + depth, + nodes_with_degree, + use_id_to_match, + v_cache, + ) + else: + # Process each edge + flat_rel, prior_edge_str_len = self._process_edge( + item, flat_rel, raw_flat_rel, i, use_id_to_match, e_cache + ) + + return flat_rel, nodes_with_degree + + def _update_vertex_degree_list( + self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str] + ) -> None: + for depth, node_str in enumerate(nodes_with_degree): + if depth >= len(vertex_degree_list): + vertex_degree_list.append(set()) + vertex_degree_list[depth].add(node_str) + + def _format_graph_query_result( + self, query_paths + ) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: + use_id_to_match = self._prop_to_match is None + subgraph = set() + subgraph_with_degree = {} + vertex_degree_list: List[Set[str]] = [] + v_cache: Set[str] = set() + e_cache: Set[Tuple[str, str, str]] = set() + + for path in query_paths: + # 1. Process each path + path_str, vertex_with_degree = self._process_path( + path, use_id_to_match, v_cache, e_cache + ) + subgraph.add(path_str) + subgraph_with_degree[path_str] = vertex_with_degree + # 2. Update vertex degree list + self._update_vertex_degree_list(vertex_degree_list, vertex_with_degree) + + return subgraph, vertex_degree_list, subgraph_with_degree + + def _get_graph_schema(self, refresh: bool = False) -> str: + if self._schema and not refresh: + return self._schema + + schema = self._client.schema() + vertex_schema = schema.getVertexLabels() + edge_schema = schema.getEdgeLabels() + relationships = schema.getRelations() + + self._schema = ( + f"Vertex properties: {vertex_schema}\n" + f"Edge properties: {edge_schema}\n" + f"Relationships: {relationships}\n" + ) + log.debug("Link(Relation): %s", relationships) + return self._schema + + @staticmethod + def _extract_label_names( + source: str, head: str = "name: ", tail: str = ", " + ) -> List[str]: + result = [] + for s in source.split(head): + end = s.find(tail) + label = s[:end] + if label: + result.append(label) + return result + + def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: + schema = self._get_graph_schema() + vertex_props_str, edge_props_str = schema.split("\n")[:2] + # TODO: rename to vertex (also need update in the schema) + vertex_props_str = ( + vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") + ) + edge_props_str = ( + edge_props_str[len("Edge properties: ") :].strip("[").strip("]") + ) + vertex_labels = self._extract_label_names(vertex_props_str) + edge_labels = self._extract_label_names(edge_props_str) + return vertex_labels, edge_labels + + def _format_graph_from_vertex(self, query_result: List[Any]) -> Set[str]: + knowledge = set() + for item in query_result: + props_str = ", ".join(f"{k}: {v}" for k, v in item["properties"].items()) + node_str = f"{item['id']}{{{props_str}}}" + knowledge.add(node_str) + return knowledge + + def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: + # 1. Extract params from context + matched_vids = context.get("match_vids") + if isinstance(context.get("max_deep"), int): + self._max_deep = context["max_deep"] + if isinstance(context.get("max_items"), int): + self._max_items = context["max_items"] + if isinstance(context.get("prop_to_match"), str): + self._prop_to_match = context["prop_to_match"] + + # 2. Extract edge_labels from graph schema + _, edge_labels = self._extract_labels_from_schema() + edge_labels_str = ",".join("'" + label + "'" for label in edge_labels) + # TODO: enhance the limit logic later + edge_limit_amount = len(edge_labels) * huge_settings.edge_limit_pre_label + + use_id_to_match = self._prop_to_match is None + if use_id_to_match: + if not matched_vids: + return context + + gremlin_query = VERTEX_QUERY_TPL.format(keywords=matched_vids) + vertexes = self._client.gremlin().exec(gremlin=gremlin_query)["data"] + log.debug("Vids gremlin query: %s", gremlin_query) + + vertex_knowledge = self._format_graph_from_vertex(query_result=vertexes) + paths: List[Any] = [] + # TODO: use generator or asyncio to speed up the query logic + for matched_vid in matched_vids: + gremlin_query = VID_QUERY_NEIGHBOR_TPL.format( + keywords=f"'{matched_vid}'", + max_deep=self._max_deep, + edge_labels=edge_labels_str, + edge_limit=edge_limit_amount, + max_items=self._max_items, + ) + log.debug("Kneighbor gremlin query: %s", gremlin_query) + paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) + + graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( + self._format_graph_query_result(query_paths=paths) + ) + + # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) + if not graph_chain_knowledge: + graph_chain_knowledge.update(vertex_knowledge) + if vertex_degree_list: + vertex_degree_list[0].update(vertex_knowledge) + else: + vertex_degree_list.append(vertex_knowledge) + else: + # WARN: When will the query enter here? + keywords = context.get("keywords") + assert keywords, "No related property(keywords) for graph query." + keywords_str = ",".join("'" + kw + "'" for kw in keywords) + gremlin_query = PROPERTY_QUERY_NEIGHBOR_TPL.format( + prop=self._prop_to_match, + keywords=keywords_str, + edge_labels=edge_labels_str, + edge_limit=edge_limit_amount, + max_deep=self._max_deep, + max_items=self._max_items, + ) + log.warning( + "Unable to find vid, downgraded to property query, please confirm if it meets expectation." + ) + + paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)[ + "data" + ] + graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( + self._format_graph_query_result(query_paths=paths) + ) + + context["graph_result"] = list(graph_chain_knowledge) + if context["graph_result"]: + context["graph_result_flag"] = 0 + context["vertex_degree_list"] = [ + list(vertex_degree) for vertex_degree in vertex_degree_list + ] + context["knowledge_with_degree"] = knowledge_with_degree + context["graph_context_head"] = ( + f"The following are graph knowledge in {self._max_deep} depth, e.g:\n" + "`vertexA--[links]-->vertexB<--[links]--vertexC ...`" + "extracted based on key entities as subject:\n" + ) + return context + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: """ Execute the graph query operation. @@ -79,8 +467,23 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: return data_json # Execute the graph query (assuming schema and semantic query have been completed in previous nodes) - graph_result = self.graph_rag_query.run(data_json) - data_json.update(graph_result) + self.init_client(data_json) + + # initial flag: -1 means no result, 0 means subgraph query, 1 means gremlin query + data_json["graph_result_flag"] = -1 + # 1. Try to perform a query based on the generated gremlin + if self._num_gremlin_generate_example >= 0: + data_json = self._gremlin_generate_query(data_json) + # 2. Try to perform a query based on subgraph-search if the previous query failed + if not data_json.get("graph_result"): + data_json = self._subgraph_query(data_json) + + if data_json.get("graph_result"): + log.debug( + "Knowledge from Graph:\n%s", "\n".join(data_json["graph_result"]) + ) + else: + log.debug("No Knowledge Extracted from Graph") log.info( f"Graph query completed, found {len(data_json.get('graph_result', []))} results" 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 3face9d63..4210b8d4a 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py @@ -15,6 +15,7 @@ import json +from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.common_op.check_schema import CheckSchema from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager @@ -46,15 +47,16 @@ def _import_schema( raise ValueError("No input data / invalid schema type") def node_init(self): - self.schema = self.wk_input.schema - self.schema = self.schema.strip() + if self.wk_input.schema is None: + return CStatus(-1, "Schema message is required in SchemaNode") + self.schema = self.wk_input.schema.strip() if self.schema.startswith("{"): try: schema = json.loads(self.schema) self.check_schema = self._import_schema(from_user_defined=schema) except json.JSONDecodeError as exc: log.error("Invalid JSON format in schema. Please check it again.") - raise ValueError("Invalid JSON format in schema.") from exc + 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) @@ -63,11 +65,7 @@ def node_init(self): def operator_schedule(self, data_json): log.debug("SchemaNode input state: %s", data_json) if self.schema.startswith("{"): - try: - return self.check_schema.run(data_json) - except json.JSONDecodeError as exc: - log.error("Invalid JSON format in schema. Please check it again.") - raise ValueError("Invalid JSON format in schema.") from exc + return self.check_schema.run(data_json) else: log.info("Get schema '%s' from graphdb.", self.schema) return self.schema_manager.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py new file mode 100644 index 000000000..2237bfabb --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py @@ -0,0 +1,43 @@ +# 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 PyCGraph import CStatus + +from hugegraph_llm.config import llm_settings +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.build_gremlin_example_index import ( + BuildGremlinExampleIndex, +) +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class BuildGremlinExampleIndexNode(BaseNode): + build_gremlin_example_index_op: BuildGremlinExampleIndex + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + if self.wk_input.examples is not None: + examples = self.wk_input.examples + else: + return CStatus(-1, "examples is required in BuildGremlinExampleIndexNode") + self.build_gremlin_example_index_op = BuildGremlinExampleIndex( + get_embedding(llm_settings), examples + ) + return super().node_init() + + def operator_schedule(self, data_json): + return self.build_gremlin_example_index_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py index bf605aa49..18e480b12 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -39,13 +39,25 @@ def node_init(self): return CStatus(-1, "graph_name is required in wk_input") embedding = get_embedding(llm_settings) - by = self.wk_input.semantic_by or "keywords" + by = ( + self.wk_input.semantic_by + if self.wk_input.semantic_by is not None + else "keywords" + ) topk_per_keyword = ( - self.wk_input.topk_per_keyword or huge_settings.topk_per_keyword + self.wk_input.topk_per_keyword + if self.wk_input.topk_per_keyword is not None + else huge_settings.topk_per_keyword + ) + topk_per_query = ( + self.wk_input.topk_per_query + if self.wk_input.topk_per_query is not None + else 10 ) - topk_per_query = self.wk_input.topk_per_query or 10 vector_dis_threshold = ( - self.wk_input.vector_dis_threshold or huge_settings.vector_dis_threshold + self.wk_input.vector_dis_threshold + if self.wk_input.vector_dis_threshold is not None + else huge_settings.vector_dis_threshold ) # Initialize the semantic ID query operator diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py index a36831526..0904b9920 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py @@ -18,7 +18,6 @@ import json from typing import Any, Dict, Optional -from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize @@ -27,13 +26,14 @@ def _stable_schema_string(state_json: Dict[str, Any]) -> str: - if "simple_schema" in state_json and state_json["simple_schema"] is not None: - return json.dumps( - state_json["simple_schema"], ensure_ascii=False, sort_keys=True - ) - if "schema" in state_json and state_json["schema"] is not None: - return json.dumps(state_json["schema"], ensure_ascii=False, sort_keys=True) - return "" + val = state_json.get("simple_schema") + if val is None: + val = state_json.get("schema") + if val is None: + return "" + if isinstance(val, str): + return val + return json.dumps(val, ensure_ascii=False, sort_keys=True) class Text2GremlinNode(BaseNode): @@ -56,7 +56,7 @@ def node_init(self): vertices=None, gremlin_prompt=gremlin_prompt, ) - return CStatus() + return super().node_init() def operator_schedule(self, data_json: Dict[str, Any]): # Ensure query exists in context; return empty if not provided diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/util.py b/hugegraph-llm/src/hugegraph_llm/nodes/util.py index 60bdc2e86..c98ad5e49 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/util.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/util.py @@ -13,10 +13,24 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any + from PyCGraph import CStatus -def init_context(obj) -> CStatus: +def init_context(obj: Any) -> CStatus: + """ + Initialize workflow context for a node. + + Retrieves wkflow_state and wkflow_input from obj's global parameters + and assigns them to obj.context and obj.wk_input respectively. + + Args: + obj: Node object with getGParamWithNoEmpty method + + Returns: + CStatus: Empty status on success, error status with code -1 on failure + """ try: obj.context = obj.getGParamWithNoEmpty("wkflow_state") obj.wk_input = obj.getGParamWithNoEmpty("wkflow_input") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py deleted file mode 100644 index 70f3d27d2..000000000 --- a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py +++ /dev/null @@ -1,81 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -from typing import Optional, List - -from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.models.llms.base import BaseLLM -from hugegraph_llm.operators.common_op.check_schema import CheckSchema -from hugegraph_llm.operators.common_op.print_result import PrintResult -from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager -from hugegraph_llm.operators.index_op.build_gremlin_example_index import BuildGremlinExampleIndex -from hugegraph_llm.operators.index_op.gremlin_example_index_query import GremlinExampleIndexQuery -from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize -from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm - - -class GremlinGenerator: - def __init__(self, llm: BaseLLM, embedding: BaseEmbedding): - self.embedding = [] - self.llm = llm - self.embedding = embedding - self.result = None - self.operators = [] - - def clear(self): - self.operators = [] - return self - - def example_index_build(self, examples): - self.operators.append(BuildGremlinExampleIndex(self.embedding, examples)) - return self - - def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): - if from_hugegraph: - self.operators.append(SchemaManager(from_hugegraph)) - elif from_user_defined: - self.operators.append(CheckSchema(from_user_defined)) - elif from_extraction: - raise NotImplementedError("Not implemented yet") - else: - raise ValueError("No input data / invalid schema type") - return self - - def example_index_query(self, num_examples): - self.operators.append(GremlinExampleIndexQuery(self.embedding, num_examples)) - return self - - def gremlin_generate_synthesize( - self, schema, gremlin_prompt: Optional[str] = None, vertices: Optional[List[str]] = None - ): - self.operators.append(GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt)) - return self - - def print_result(self): - self.operators.append(PrintResult()) - return self - - @log_time("total time") - @record_rpm - def run(self, **kwargs): - context = kwargs - for operator in self.operators: - context = self._run_operator(operator, context) - return context - - @log_operator_time - def _run_operator(self, operator, context): - return operator.run(context) 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 52626b72b..ba4392f7c 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 @@ -40,7 +40,6 @@ def run(self, data: dict) -> Dict[str, Any]: schema = data.get("schema") vertices = data.get("vertices", []) edges = data.get("edges", []) - print(f"get schema {schema}") if not vertices and not edges: log.critical( "(Loading) Both vertices and edges are empty. Please check the input data again." @@ -50,7 +49,9 @@ def run(self, data: dict) -> Dict[str, Any]: if not schema: # TODO: ensure the function works correctly (update the logic later) self.schema_free_mode(data.get("triples", [])) - log.warning("Using schema_free mode, could try schema_define mode for better effect!") + log.warning( + "Using schema_free mode, could try schema_define mode for better effect!" + ) else: self.init_schema_if_need(schema) self.load_into_graph(vertices, edges, schema) @@ -66,7 +67,9 @@ def _set_default_property(self, key, input_properties, property_label_map): # list or set default_value = [] input_properties[key] = default_value - log.warning("Property '%s' missing in vertex, set to '%s' for now", key, default_value) + log.warning( + "Property '%s' missing in vertex, set to '%s' for now", key, default_value + ) def _handle_graph_creation(self, func, *args, **kwargs): try: @@ -80,9 +83,13 @@ 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) - vertex_label_map = {v_label["name"]: v_label for v_label in schema["vertexlabels"]} + 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"]} + property_label_map = { + p_label["name"]: p_label for p_label in schema["propertykeys"] + } for vertex in vertices: input_label = vertex["label"] @@ -98,7 +105,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- vertex_label = vertex_label_map[input_label] primary_keys = vertex_label["primary_keys"] nullable_keys = vertex_label.get("nullable_keys", []) - non_null_keys = [key for key in vertex_label["properties"] if key not in nullable_keys] + non_null_keys = [ + key for key in vertex_label["properties"] if key not in nullable_keys + ] has_problem = False # 2. Handle primary-keys mode vertex @@ -130,7 +139,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- # 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) + 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(): @@ -167,7 +178,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- 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) + self._handle_graph_creation( + self.client.graph().addEdge, label, start, end, properties + ) def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -191,18 +204,20 @@ def init_schema_if_need(self, schema: dict): source_vertex_label = edge["source_label"] target_vertex_label = edge["target_label"] properties = edge["properties"] - self.schema.edgeLabel(edge_label).sourceLabel(source_vertex_label).targetLabel( - target_vertex_label - ).properties(*properties).nullableKeys(*properties).ifNotExist().create() + self.schema.edgeLabel(edge_label).sourceLabel( + source_vertex_label + ).targetLabel(target_vertex_label).properties(*properties).nullableKeys( + *properties + ).ifNotExist().create() def schema_free_mode(self, 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( - "name" - ).ifNotExist().create() + self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel( + "vertex" + ).properties("name").ifNotExist().create() self.schema.indexLabel("vertexByName").onV("vertex").by( "name" @@ -262,7 +277,9 @@ def _set_property_data_type(self, property_key, data_type): log.warning("UUID type is not supported, use text instead") property_key.asText() else: - log.error("Unknown data type %s for property_key %s", data_type, property_key) + log.error( + "Unknown data type %s for property_key %s", data_type, property_key + ) def _set_property_cardinality(self, property_key, cardinality): if cardinality == PropertyCardinality.SINGLE: @@ -272,9 +289,13 @@ def _set_property_cardinality(self, property_key, cardinality): elif cardinality == PropertyCardinality.SET: property_key.valueSet() else: - log.error("Unknown cardinality %s for property_key %s", cardinality, property_key) + 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: + def _check_property_data_type( + self, data_type: str, cardinality: str, value + ) -> bool: if cardinality in ( PropertyCardinality.LIST.value, PropertyCardinality.SET.value, @@ -304,7 +325,9 @@ def _check_single_data_type(self, data_type: str, value) -> bool: if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below - if data_type == PropertyDataType.DATE.value: # the format should be "yyyy-MM-dd" + if ( + data_type == PropertyDataType.DATE.value + ): # the format should be "yyyy-MM-dd" import re return isinstance(value, str) and re.match(r"^\d{4}-\d{2}-\d{2}$", value) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py deleted file mode 100644 index bcff5f07b..000000000 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py +++ /dev/null @@ -1,455 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -import json -from typing import Any, Dict, Optional, List, Set, Tuple - -from hugegraph_llm.config import huge_settings, prompt -from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.models.llms.base import BaseLLM -from hugegraph_llm.operators.gremlin_generate_task import GremlinGenerator -from hugegraph_llm.utils.log import log -from pyhugegraph.client import PyHugeClient - -# TODO: remove 'as('subj)' step -VERTEX_QUERY_TPL = "g.V({keywords}).limit(8).as('subj').toList()" - -# TODO: we could use a simpler query (like kneighbor-api to get the edges) -# TODO: test with profile()/explain() to speed up the query -VID_QUERY_NEIGHBOR_TPL = """\ -g.V({keywords}) -.repeat( - bothE({edge_labels}).limit({edge_limit}).otherV().dedup() -).times({max_deep}).emit() -.simplePath() -.path() -.by(project('label', 'id', 'props') - .by(label()) - .by(id()) - .by(valueMap().by(unfold())) -) -.by(project('label', 'inV', 'outV', 'props') - .by(label()) - .by(inV().id()) - .by(outV().id()) - .by(valueMap().by(unfold())) -) -.limit({max_items}) -.toList() -""" - -PROPERTY_QUERY_NEIGHBOR_TPL = """\ -g.V().has('{prop}', within({keywords})) -.repeat( - bothE({edge_labels}).limit({edge_limit}).otherV().dedup() -).times({max_deep}).emit() -.simplePath() -.path() -.by(project('label', 'props') - .by(label()) - .by(valueMap().by(unfold())) -) -.by(project('label', 'inV', 'outV', 'props') - .by(label()) - .by(inV().values('{prop}')) - .by(outV().values('{prop}')) - .by(valueMap().by(unfold())) -) -.limit({max_items}) -.toList() -""" - - -class GraphRAGQuery: - def __init__( - self, - max_deep: int = 2, - max_graph_items: int = huge_settings.max_graph_items, - prop_to_match: Optional[str] = None, - llm: Optional[BaseLLM] = None, - embedding: Optional[BaseEmbedding] = None, - max_v_prop_len: Optional[int] = 2048, - max_e_prop_len: Optional[int] = 256, - num_gremlin_generate_example: Optional[int] = -1, - gremlin_prompt: Optional[str] = None, - ): - 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, - ) - self._max_deep = max_deep - self._max_items = max_graph_items - self._prop_to_match = prop_to_match - self._schema = "" - self._limit_property = huge_settings.limit_property.lower() == "true" - self._max_v_prop_len = max_v_prop_len - self._max_e_prop_len = max_e_prop_len - self._gremlin_generator = GremlinGenerator( - llm=llm, - embedding=embedding, - ) - self._num_gremlin_generate_example = num_gremlin_generate_example - self._gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt - - def run(self, context: Dict[str, Any]) -> Dict[str, Any]: - self.init_client(context) - - # initial flag: -1 means no result, 0 means subgraph query, 1 means gremlin query - context["graph_result_flag"] = -1 - # 1. Try to perform a query based on the generated gremlin - if self._num_gremlin_generate_example >= 0: - context = self._gremlin_generate_query(context) - # 2. Try to perform a query based on subgraph-search if the previous query failed - if not context.get("graph_result"): - context = self._subgraph_query(context) - - if context.get("graph_result"): - log.debug("Knowledge from Graph:\n%s", "\n".join(context["graph_result"])) - else: - log.debug("No Knowledge Extracted from Graph") - return context - - def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: - query = context["query"] - vertices = context.get("match_vids") - query_embedding = context.get("query_embedding") - - self._gremlin_generator.clear() - self._gremlin_generator.example_index_query(num_examples=self._num_gremlin_generate_example) - gremlin_response = self._gremlin_generator.gremlin_generate_synthesize( - context["simple_schema"], vertices=vertices, gremlin_prompt=self._gremlin_prompt - ).run(query=query, query_embedding=query_embedding) - if self._num_gremlin_generate_example > 0: - gremlin = gremlin_response["result"] - else: - gremlin = gremlin_response["raw_result"] - log.info("Generated gremlin: %s", gremlin) - context["gremlin"] = gremlin - try: - result = self._client.gremlin().exec(gremlin=gremlin)["data"] - if result == [None]: - result = [] - context["graph_result"] = [json.dumps(item, ensure_ascii=False) for item in result] - if context["graph_result"]: - context["graph_result_flag"] = 1 - context["graph_context_head"] = ( - f"The following are graph query result " f"from gremlin query `{gremlin}`.\n" - ) - except Exception as e: # pylint: disable=broad-except - log.error(e) - context["graph_result"] = "" - return context - - def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: - # 1. Extract params from context - matched_vids = context.get("match_vids") - if isinstance(context.get("max_deep"), int): - self._max_deep = context["max_deep"] - if isinstance(context.get("max_items"), int): - self._max_items = context["max_items"] - if isinstance(context.get("prop_to_match"), str): - self._prop_to_match = context["prop_to_match"] - - # 2. Extract edge_labels from graph schema - _, edge_labels = self._extract_labels_from_schema() - edge_labels_str = ",".join("'" + label + "'" for label in edge_labels) - # TODO: enhance the limit logic later - edge_limit_amount = len(edge_labels) * huge_settings.edge_limit_pre_label - - use_id_to_match = self._prop_to_match is None - if use_id_to_match: - if not matched_vids: - return context - - gremlin_query = VERTEX_QUERY_TPL.format(keywords=matched_vids) - vertexes = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - log.debug("Vids gremlin query: %s", gremlin_query) - - vertex_knowledge = self._format_graph_from_vertex(query_result=vertexes) - paths: List[Any] = [] - # TODO: use generator or asyncio to speed up the query logic - for matched_vid in matched_vids: - gremlin_query = VID_QUERY_NEIGHBOR_TPL.format( - keywords=f"'{matched_vid}'", - max_deep=self._max_deep, - edge_labels=edge_labels_str, - edge_limit=edge_limit_amount, - max_items=self._max_items, - ) - log.debug("Kneighbor gremlin query: %s", gremlin_query) - paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) - - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( - self._format_graph_query_result(query_paths=paths) - ) - - # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) - if not graph_chain_knowledge: - graph_chain_knowledge.update(vertex_knowledge) - if vertex_degree_list: - vertex_degree_list[0].update(vertex_knowledge) - else: - vertex_degree_list.append(vertex_knowledge) - else: - # WARN: When will the query enter here? - keywords = context.get("keywords") - assert keywords, "No related property(keywords) for graph query." - keywords_str = ",".join("'" + kw + "'" for kw in keywords) - gremlin_query = PROPERTY_QUERY_NEIGHBOR_TPL.format( - prop=self._prop_to_match, - keywords=keywords_str, - edge_labels=edge_labels_str, - edge_limit=edge_limit_amount, - max_deep=self._max_deep, - max_items=self._max_items, - ) - log.warning( - "Unable to find vid, downgraded to property query, please confirm if it meets expectation." - ) - - paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( - self._format_graph_query_result(query_paths=paths) - ) - - context["graph_result"] = list(graph_chain_knowledge) - if context["graph_result"]: - context["graph_result_flag"] = 0 - context["vertex_degree_list"] = [ - list(vertex_degree) for vertex_degree in vertex_degree_list - ] - context["knowledge_with_degree"] = knowledge_with_degree - context["graph_context_head"] = ( - f"The following are graph knowledge in {self._max_deep} depth, e.g:\n" - "`vertexA--[links]-->vertexB<--[links]--vertexC ...`" - "extracted based on key entities as subject:\n" - ) - return context - - # TODO: move this method to a util file for reuse (remove self param) - def init_client(self, context): - """Initialize the HugeGraph client from context or default settings.""" - # pylint: disable=R0915 (too-many-statements) - if self._client is None: - if isinstance(context.get("graph_client"), PyHugeClient): - self._client = context["graph_client"] - else: - url = context.get("url") or "http://localhost:8080" - graph = context.get("graph") or "hugegraph" - user = context.get("user") or "admin" - pwd = context.get("pwd") or "admin" - gs = context.get("graphspace") or None - self._client = PyHugeClient(url, graph, user, pwd, gs) - assert self._client is not None, "No valid graph to search." - - def get_vertex_details(self, vertex_ids: List[str]) -> List[Dict[str, Any]]: - if not vertex_ids: - return [] - - formatted_ids = ", ".join(f"'{vid}'" for vid in vertex_ids) - gremlin_query = f"g.V({formatted_ids}).limit(20)" - result = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - return result - - def _format_graph_from_vertex(self, query_result: List[Any]) -> Set[str]: - knowledge = set() - for item in query_result: - props_str = ", ".join(f"{k}: {v}" for k, v in item["properties"].items()) - node_str = f"{item['id']}{{{props_str}}}" - knowledge.add(node_str) - return knowledge - - def _format_graph_query_result( - self, query_paths - ) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: - use_id_to_match = self._prop_to_match is None - subgraph = set() - subgraph_with_degree = {} - vertex_degree_list: List[Set[str]] = [] - v_cache: Set[str] = set() - e_cache: Set[Tuple[str, str, str]] = set() - - for path in query_paths: - # 1. Process each path - path_str, vertex_with_degree = self._process_path( - path, use_id_to_match, v_cache, e_cache - ) - subgraph.add(path_str) - subgraph_with_degree[path_str] = vertex_with_degree - # 2. Update vertex degree list - self._update_vertex_degree_list(vertex_degree_list, vertex_with_degree) - - return subgraph, vertex_degree_list, subgraph_with_degree - - def _process_path( - self, - path: Any, - use_id_to_match: bool, - v_cache: Set[str], - e_cache: Set[Tuple[str, str, str]], - ) -> Tuple[str, List[str]]: - flat_rel = "" - raw_flat_rel = path["objects"] - assert len(raw_flat_rel) % 2 == 1, "The length of raw_flat_rel should be odd." - - node_cache = set() - prior_edge_str_len = 0 - depth = 0 - nodes_with_degree = [] - - for i, item in enumerate(raw_flat_rel): - if i % 2 == 0: - # Process each vertex - flat_rel, prior_edge_str_len, depth = self._process_vertex( - item, - flat_rel, - node_cache, - prior_edge_str_len, - depth, - nodes_with_degree, - use_id_to_match, - v_cache, - ) - else: - # Process each edge - flat_rel, prior_edge_str_len = self._process_edge( - item, flat_rel, raw_flat_rel, i, use_id_to_match, e_cache - ) - - return flat_rel, nodes_with_degree - - def _process_vertex( - self, - item: Any, - flat_rel: str, - node_cache: Set[str], - prior_edge_str_len: int, - depth: int, - nodes_with_degree: List[str], - use_id_to_match: bool, - v_cache: Set[str], - ) -> Tuple[str, int, int]: - matched_str = item["id"] if use_id_to_match else item["props"][self._prop_to_match] - if matched_str in node_cache: - flat_rel = flat_rel[:-prior_edge_str_len] - return flat_rel, prior_edge_str_len, depth - - node_cache.add(matched_str) - props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v - ) - - # TODO: we may remove label id or replace with label name - if matched_str in v_cache: - node_str = matched_str - else: - v_cache.add(matched_str) - node_str = f"{item['id']}{{{props_str}}}" - - flat_rel += node_str - nodes_with_degree.append(node_str) - depth += 1 - return flat_rel, prior_edge_str_len, depth - - def _process_edge( - self, - item: Any, - path_str: str, - raw_flat_rel: List[Any], - i: int, - use_id_to_match: bool, - e_cache: Set[Tuple[str, str, str]], - ) -> Tuple[str, int]: - props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v - ) - props_str = f"{{{props_str}}}" if props_str else "" - prev_matched_str = ( - raw_flat_rel[i - 1]["id"] - if use_id_to_match - else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] - ) - - edge_key = (item["inV"], item["label"], item["outV"]) - if edge_key not in e_cache: - e_cache.add(edge_key) - edge_label = f"{item['label']}{props_str}" - else: - edge_label = item["label"] - - edge_str = ( - f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" - ) - path_str += edge_str - prior_edge_str_len = len(edge_str) - return path_str, prior_edge_str_len - - def _update_vertex_degree_list( - self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str] - ) -> None: - for depth, node_str in enumerate(nodes_with_degree): - if depth >= len(vertex_degree_list): - vertex_degree_list.append(set()) - vertex_degree_list[depth].add(node_str) - - def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: - schema = self._get_graph_schema() - vertex_props_str, edge_props_str = schema.split("\n")[:2] - # TODO: rename to vertex (also need update in the schema) - vertex_props_str = vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") - edge_props_str = edge_props_str[len("Edge properties: ") :].strip("[").strip("]") - vertex_labels = self._extract_label_names(vertex_props_str) - edge_labels = self._extract_label_names(edge_props_str) - return vertex_labels, edge_labels - - @staticmethod - def _extract_label_names(source: str, head: str = "name: ", tail: str = ", ") -> List[str]: - result = [] - for s in source.split(head): - end = s.find(tail) - label = s[:end] - if label: - result.append(label) - return result - - def _get_graph_schema(self, refresh: bool = False) -> str: - if self._schema and not refresh: - return self._schema - - schema = self._client.schema() - vertex_schema = schema.getVertexLabels() - edge_schema = schema.getEdgeLabels() - relationships = schema.getRelations() - - self._schema = ( - f"Vertex properties: {vertex_schema}\n" - f"Edge properties: {edge_schema}\n" - f"Relationships: {relationships}\n" - ) - log.debug("Link(Relation): %s", relationships) - return self._schema - - def _limit_property_query(self, value: Optional[str], item_type: str) -> Optional[str]: - # NOTE: we skip the filter for list/set type (e.g., list of string, add it if needed) - if not self._limit_property or not isinstance(value, str): - return value - - max_len = self._max_v_prop_len if item_type == "v" else self._max_e_prop_len - return value[:max_len] if value else value 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 5689a59ac..2ed4e840a 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 @@ -37,11 +37,15 @@ def __init__(self, embedding: BaseEmbedding): self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) + self.index_dir = str( + os.path.join(resource_path, self.folder_name, "graph_vids") + ) self.filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(embedding, "model_name", None) ) - self.vid_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) + self.vid_index = VectorIndex.from_index_file( + self.index_dir, self.filename_prefix + ) self.embedding = embedding self.sm = SchemaManager(huge_settings.graph_name) @@ -50,19 +54,27 @@ def _extract_names(self, vertices: list[str]) -> list[str]: def run(self, context: Dict[str, Any]) -> Dict[str, Any]: vertexlabels = self.sm.schema.getSchema()["vertexlabels"] - all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) + all_pk_flag = bool(vertexlabels) and all( + data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels + ) past_vids = self.vid_index.properties # TODO: We should build vid vector index separately, especially when the vertices may be very large - present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py + present_vids = context[ + "vertices" + ] # Warning: data truncated by fetch_graph_data.py removed_vids = set(past_vids) - set(present_vids) removed_num = self.vid_index.remove(removed_vids) added_vids = list(set(present_vids) - set(past_vids)) if added_vids: - vids_to_process = self._extract_names(added_vids) if all_pk_flag else added_vids - added_embeddings = asyncio.run(get_embeddings_parallel(self.embedding, vids_to_process)) + vids_to_process = ( + self._extract_names(added_vids) if all_pk_flag else added_vids + ) + added_embeddings = asyncio.run( + get_embeddings_parallel(self.embedding, 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.to_index_file(self.index_dir, self.filename_prefix) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py deleted file mode 100644 index 3b5c63103..000000000 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ /dev/null @@ -1,120 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - - -from typing import Dict, Any, Optional, Literal, Union, List - -from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.models.llms.base import BaseLLM -from hugegraph_llm.operators.common_op.check_schema import CheckSchema -from hugegraph_llm.operators.common_op.print_result import PrintResult -from hugegraph_llm.operators.document_op.chunk_split import ChunkSplit -from hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph import Commit2Graph -from hugegraph_llm.operators.hugegraph_op.fetch_graph_data import FetchGraphData -from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager -from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex -from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex -from hugegraph_llm.operators.llm_op.disambiguate_data import DisambiguateData -from hugegraph_llm.operators.llm_op.info_extract import InfoExtract -from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract -from hugegraph_llm.operators.llm_op.schema_build import SchemaBuilder -from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm -from pyhugegraph.client import PyHugeClient - - -class KgBuilder: - def __init__( - self, - llm: BaseLLM, - embedding: Optional[BaseEmbedding] = None, - graph: Optional[PyHugeClient] = None, - ): - self.operators = [] - self.llm = llm - self.embedding = embedding - self.graph = graph - self.result = None - - def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): - if from_hugegraph: - self.operators.append(SchemaManager(from_hugegraph)) - elif from_user_defined: - self.operators.append(CheckSchema(from_user_defined)) - elif from_extraction: - raise NotImplementedError("Not implemented yet") - else: - raise ValueError("No input data / invalid schema type") - return self - - def fetch_graph_data(self): - self.operators.append(FetchGraphData(self.graph)) - return self - - def chunk_split( - self, - text: Union[str, List[str]], # text to be split - split_type: Literal["document", "paragraph", "sentence"] = "document", - language: Literal["zh", "en"] = "zh", - ): - self.operators.append(ChunkSplit(text, split_type, language)) - return self - - def extract_info( - self, - example_prompt: Optional[str] = None, - extract_type: Literal["triples", "property_graph"] = "triples", - ): - if extract_type == "triples": - self.operators.append(InfoExtract(self.llm, example_prompt)) - elif extract_type == "property_graph": - self.operators.append(PropertyGraphExtract(self.llm, example_prompt)) - return self - - def disambiguate_word_sense(self): - self.operators.append(DisambiguateData(self.llm)) - return self - - def commit_to_hugegraph(self): - self.operators.append(Commit2Graph()) - return self - - def build_vertex_id_semantic_index(self): - self.operators.append(BuildSemanticIndex(self.embedding)) - return self - - def build_vector_index(self): - self.operators.append(BuildVectorIndex(self.embedding)) - return self - - def print_result(self): - self.operators.append(PrintResult()) - return self - - def build_schema(self): - self.operators.append(SchemaBuilder(self.llm)) - return self - - @log_time("total time") - @record_rpm - def run(self, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: - for operator in self.operators: - context = self._run_operator(operator, context) - return context - - @log_operator_time - def _run_operator(self, operator, context) -> Dict[str, Any]: - return operator.run(context) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py similarity index 58% rename from hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py rename to hugegraph-llm/src/hugegraph_llm/operators/operator_list.py index 330890b5d..46ddbd50e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py @@ -14,45 +14,131 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - - -from typing import Dict, Any, Optional, List, Literal +from typing import Optional, List, Literal, Union from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.models.llms.base import BaseLLM -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.common_op.merge_dedup_rerank import MergeDedupRerank +from hugegraph_llm.operators.common_op.check_schema import CheckSchema from hugegraph_llm.operators.common_op.print_result import PrintResult -from hugegraph_llm.operators.document_op.word_extract import WordExtract -from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager +from hugegraph_llm.operators.index_op.build_gremlin_example_index import ( + BuildGremlinExampleIndex, +) +from hugegraph_llm.operators.index_op.gremlin_example_index_query import ( + GremlinExampleIndexQuery, +) +from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize +from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm +from hugegraph_llm.operators.hugegraph_op.fetch_graph_data import FetchGraphData +from hugegraph_llm.operators.document_op.chunk_split import ChunkSplit +from hugegraph_llm.operators.llm_op.info_extract import InfoExtract +from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract +from hugegraph_llm.operators.llm_op.disambiguate_data import DisambiguateData +from hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph import Commit2Graph +from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex +from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex +from hugegraph_llm.operators.document_op.word_extract import WordExtract +from hugegraph_llm.operators.llm_op.keyword_extract import KeywordExtract from hugegraph_llm.operators.index_op.semantic_id_query import SemanticIdQuery from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery +from hugegraph_llm.operators.common_op.merge_dedup_rerank import MergeDedupRerank from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize -from hugegraph_llm.operators.llm_op.keyword_extract import KeywordExtract -from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm -from hugegraph_llm.config import prompt, huge_settings +from hugegraph_llm.config import huge_settings +from pyhugegraph.client import PyHugeClient -class RAGPipeline: - """ - RAGPipeline is a (core)class that encapsulates a series of operations for extracting information from text, - querying graph databases and vector indices, merging and re-ranking results, and generating answers. - """ +class OperatorList: + def __init__( + self, + llm: BaseLLM, + embedding: BaseEmbedding, + graph: Optional[PyHugeClient] = None, + ): + self.llm = llm + self.embedding = embedding + self.result = None + self.operators = [] + self.graph = graph - def __init__(self, llm: Optional[BaseLLM] = None, embedding: Optional[BaseEmbedding] = None): - """ - Initialize the RAGPipeline with optional LLM and embedding models. + def clear(self): + self.operators = [] + return self - :param llm: Optional LLM model to use. - :param embedding: Optional embedding model to use. - """ - self._chat_llm = llm or LLMs().get_chat_llm() - self._extract_llm = llm or LLMs().get_extract_llm() - self._text2gqlt_llm = llm or LLMs().get_text2gql_llm() - self._embedding = embedding or Embeddings().get_embedding() - self._operators: List[Any] = [] + def example_index_build(self, examples): + self.operators.append(BuildGremlinExampleIndex(self.embedding, examples)) + return self + + def import_schema( + self, from_hugegraph=None, from_extraction=None, from_user_defined=None + ): + if from_hugegraph: + self.operators.append(SchemaManager(from_hugegraph)) + elif from_user_defined: + self.operators.append(CheckSchema(from_user_defined)) + elif from_extraction: + raise NotImplementedError("Not implemented yet") + else: + raise ValueError("No input data / invalid schema type") + return self + + def example_index_query(self, num_examples): + self.operators.append(GremlinExampleIndexQuery(self.embedding, num_examples)) + return self + + def gremlin_generate_synthesize( + self, + schema, + gremlin_prompt: Optional[str] = None, + vertices: Optional[List[str]] = None, + ): + self.operators.append( + GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt) + ) + return self + + def print_result(self): + self.operators.append(PrintResult()) + return self + + def fetch_graph_data(self): + self.operators.append(FetchGraphData(self.graph)) + return self + + def chunk_split( + self, + text: Union[str, List[str]], # text to be split + split_type: Literal["document", "paragraph", "sentence"] = "document", + language: Literal["zh", "en"] = "zh", + ): + self.operators.append(ChunkSplit(text, split_type, language)) + return self + + def extract_info( + self, + example_prompt: Optional[str] = None, + extract_type: Literal["triples", "property_graph"] = "triples", + ): + if extract_type == "triples": + self.operators.append(InfoExtract(self.llm, example_prompt)) + elif extract_type == "property_graph": + self.operators.append(PropertyGraphExtract(self.llm, example_prompt)) + return self + + def disambiguate_word_sense(self): + self.operators.append(DisambiguateData(self.llm)) + return self + + def commit_to_hugegraph(self): + self.operators.append(Commit2Graph()) + return self + + def build_vertex_id_semantic_index(self): + self.operators.append(BuildSemanticIndex(self.embedding)) + return self + + def build_vector_index(self): + self.operators.append(BuildVectorIndex(self.embedding)) + return self def extract_word(self, text: Optional[str] = None, language: str = "english"): """ @@ -62,7 +148,7 @@ def extract_word(self, text: Optional[str] = None, language: str = "english"): :param language: Language of the text. :return: Self-instance for chaining. """ - self._operators.append(WordExtract(text=text, language=language)) + self.operators.append(WordExtract(text=text, language=language)) return self def extract_keywords( @@ -81,7 +167,7 @@ def extract_keywords( :param extract_template: Template for keyword extraction. :return: Self-instance for chaining. """ - self._operators.append( + self.operators.append( KeywordExtract( text=text, max_keywords=max_keywords, @@ -91,10 +177,6 @@ def extract_keywords( ) return self - def import_schema(self, graph_name: str): - self._operators.append(SchemaManager(graph_name)) - return self - def keywords_to_vid( self, by: Literal["query", "keywords"] = "keywords", @@ -110,9 +192,9 @@ def keywords_to_vid( :param vector_dis_threshold: Vector distance threshold. :return: Self-instance for chaining. """ - self._operators.append( + self.operators.append( SemanticIdQuery( - embedding=self._embedding, + embedding=self.embedding, by=by, topk_per_keyword=topk_per_keyword, topk_per_query=topk_per_query, @@ -121,41 +203,6 @@ def keywords_to_vid( ) return self - def query_graphdb( - self, - max_deep: int = 2, - max_graph_items: int = huge_settings.max_graph_items, - max_v_prop_len: int = 2048, - max_e_prop_len: int = 256, - prop_to_match: Optional[str] = None, - num_gremlin_generate_example: Optional[int] = -1, - gremlin_prompt: Optional[str] = prompt.gremlin_generate_prompt, - ): - """ - Add a graph RAG query operator to the pipeline. - - :param max_deep: Maximum depth for the graph query. - :param max_graph_items: Maximum number of items to retrieve. - :param max_v_prop_len: Maximum length of vertex properties. - :param max_e_prop_len: Maximum length of edge properties. - :param prop_to_match: Property to match in the graph. - :param num_gremlin_generate_example: Number of examples to generate. - :param gremlin_prompt: Gremlin prompt for generating examples. - :return: Self-instance for chaining. - """ - self._operators.append( - GraphRAGQuery( - max_deep=max_deep, - max_graph_items=max_graph_items, - max_v_prop_len=max_v_prop_len, - max_e_prop_len=max_e_prop_len, - prop_to_match=prop_to_match, - num_gremlin_generate_example=num_gremlin_generate_example, - gremlin_prompt=gremlin_prompt, - ) - ) - return self - def query_vector_index(self, max_items: int = 3): """ Add a vector index query operator to the pipeline. @@ -163,9 +210,9 @@ def query_vector_index(self, max_items: int = 3): :param max_items: Maximum number of items to retrieve. :return: Self-instance for chaining. """ - self._operators.append( + self.operators.append( VectorIndexQuery( - embedding=self._embedding, + embedding=self.embedding, topk=max_items, ) ) @@ -184,9 +231,9 @@ def merge_dedup_rerank( :return: Self-instance for chaining. """ - self._operators.append( + self.operators.append( MergeDedupRerank( - embedding=self._embedding, + embedding=self.embedding, graph_ratio=graph_ratio, method=rerank_method, near_neighbor_first=near_neighbor_first, @@ -214,7 +261,7 @@ def synthesize_answer( :param answer_prompt: Template for the answer synthesis prompt. :return: Self-instance for chaining. """ - self._operators.append( + self.operators.append( AnswerSynthesize( raw_answer=raw_answer, vector_only_answer=vector_only_answer, @@ -225,32 +272,11 @@ def synthesize_answer( ) return self - def print_result(self): - """ - Add a print result operator to the pipeline. - - :return: Self-instance for chaining. - """ - self._operators.append(PrintResult()) - return self - @log_time("total time") @record_rpm - def run(self, **kwargs) -> Dict[str, Any]: - """ - Execute all operators in the pipeline in sequence. - - :param kwargs: Additional context to pass to operators. - :return: Final context after all operators have been executed. - """ - if len(self._operators) == 0: - self.extract_keywords().query_graphdb( - max_graph_items=kwargs.get("max_graph_items") - ).synthesize_answer() - + def run(self, **kwargs): context = kwargs - - for operator in self._operators: + for operator in self.operators: context = self._run_operator(operator, context) return context diff --git a/hugegraph-llm/src/hugegraph_llm/operators/util.py b/hugegraph-llm/src/hugegraph_llm/operators/util.py deleted file mode 100644 index 60bdc2e86..000000000 --- a/hugegraph-llm/src/hugegraph_llm/operators/util.py +++ /dev/null @@ -1,27 +0,0 @@ -# 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 PyCGraph import CStatus - - -def init_context(obj) -> CStatus: - try: - obj.context = obj.getGParamWithNoEmpty("wkflow_state") - obj.wk_input = obj.getGParamWithNoEmpty("wkflow_input") - if obj.context is None or obj.wk_input is None: - return CStatus(-1, "Required workflow parameters not found") - return CStatus() - except Exception as e: - return CStatus(-1, f"Failed to initialize context: {str(e)}") diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 3a6fd3c1c..9c59a42e6 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -15,62 +15,70 @@ from PyCGraph import GParam, CStatus -from typing import Union, List, Optional, Any +from typing import Union, List, Optional, Any, Dict class WkFlowInput(GParam): - texts: Union[str, List[str]] = None # texts input used by ChunkSplit Node - language: str = None # language configuration used by ChunkSplit Node - split_type: str = None # split type used by ChunkSplit Node - example_prompt: str = None # need by graph information extract - schema: str = None # Schema information requeired by SchemaNode - data_json = None - extract_type = None - query_examples = None - few_shot_schema = None + texts: Optional[Union[str, List[str]]] = None # texts input used by ChunkSplit Node + 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 + schema: Optional[str] = None # Schema information requeired by SchemaNode + graph_name: Optional[str] = None # used by SchemaManager + data_json: Optional[dict] = None + extract_type: Optional[str] = None + query_examples: Optional[Any] = None + few_shot_schema: Optional[Any] = None # Fields related to PromptGenerate - source_text: str = None # Original text - scenario: str = None # Scenario description - example_name: str = None # Example name + source_text: Optional[str] = None # Original text + scenario: Optional[str] = None # Scenario description + example_name: Optional[str] = None # Example name # Fields for Text2Gremlin - example_num: int = None - gremlin_prompt: str = None + example_num: Optional[int] = None + gremlin_prompt: Optional[str] = None requested_outputs: Optional[List[str]] = None # RAG Flow related fields - query: str = None # User query for RAG - vector_search: bool = None # Enable vector search - graph_search: bool = None # Enable graph search - raw_answer: bool = None # Return raw answer - vector_only_answer: bool = None # Vector only answer mode - graph_only_answer: bool = None # Graph only answer mode - graph_vector_answer: bool = None # Combined graph and vector answer - graph_ratio: float = None # Graph ratio for merging - rerank_method: str = None # Reranking method - near_neighbor_first: bool = None # Near neighbor first flag - custom_related_information: str = None # Custom related information - answer_prompt: str = None # Answer generation prompt - keywords_extract_prompt: str = None # Keywords extraction prompt - gremlin_tmpl_num: int = None # Gremlin template number - gremlin_prompt: str = None # Gremlin generation prompt - max_graph_items: int = None # Maximum graph items - topk_return_results: int = None # Top-k return results - vector_dis_threshold: float = None # Vector distance threshold - topk_per_keyword: int = None # Top-k per keyword - max_keywords: int = None - max_items: int = None + query: Optional[str] = None # User query for RAG + vector_search: Optional[bool] = None # Enable vector search + graph_search: Optional[bool] = None # Enable graph search + raw_answer: Optional[bool] = None # Return raw answer + vector_only_answer: Optional[bool] = None # Vector only answer mode + graph_only_answer: Optional[bool] = None # Graph only answer mode + graph_vector_answer: Optional[bool] = None # Combined graph and vector answer + graph_ratio: Optional[float] = None # Graph ratio for merging + rerank_method: Optional[str] = None # Reranking method + near_neighbor_first: Optional[bool] = None # Near neighbor first flag + custom_related_information: Optional[str] = None # Custom related information + answer_prompt: Optional[str] = None # Answer generation prompt + keywords_extract_prompt: Optional[str] = None # Keywords extraction prompt + gremlin_tmpl_num: Optional[int] = None # Gremlin template number + gremlin_prompt: Optional[str] = None # Gremlin generation prompt + max_graph_items: Optional[int] = None # Maximum graph items + topk_return_results: Optional[int] = None # Top-k return results + vector_dis_threshold: Optional[float] = None # Vector distance threshold + topk_per_keyword: Optional[int] = None # Top-k per keyword + max_keywords: Optional[int] = None + max_items: Optional[int] = None # Semantic query related fields - semantic_by: str = None # Semantic query method - topk_per_query: int = None # Top-k per query + semantic_by: Optional[str] = None # Semantic query method + topk_per_query: Optional[int] = None # Top-k per query # Graph query related fields - max_deep: int = None # Maximum depth for graph traversal - max_v_prop_len: int = None # Maximum vertex property length - max_e_prop_len: int = None # Maximum edge property length - prop_to_match: str = None # Property to match + max_deep: Optional[int] = None # Maximum depth for graph traversal + max_v_prop_len: Optional[int] = None # Maximum vertex property length + max_e_prop_len: Optional[int] = None # Maximum edge property length + prop_to_match: Optional[str] = None # Property to match - stream: bool = None # used for recognize stream mode + stream: Optional[bool] = None # used for recognize stream mode + + # used for rag_recall api + is_graph_rag_recall: bool = False + is_vector_only: bool = False + + # used for build text2gremin index + examples: Optional[List[Dict[str, str]]] = None def reset(self, _: CStatus) -> None: self.texts = None @@ -123,6 +131,8 @@ def reset(self, _: CStatus) -> None: self.prop_to_match = None self.stream = None + self.examples = None + class WkFlowState(GParam): schema: Optional[str] = None # schema message @@ -134,9 +144,9 @@ class WkFlowState(GParam): call_count: Optional[int] = None keywords: Optional[List[str]] = None - vector_result = None - graph_result = None - keywords_embeddings = None + vector_result: Optional[Any] = None + graph_result: Optional[Any] = None + keywords_embeddings: Optional[Any] = None generated_extract_prompt: Optional[str] = None # Fields for Text2Gremlin results @@ -146,16 +156,14 @@ class WkFlowState(GParam): template_exec_res: Optional[Any] = None raw_exec_res: Optional[Any] = None - match_vids = None - vector_result = None - graph_result = None + match_vids: Optional[Any] = None - raw_answer: str = None - vector_only_answer: str = None - graph_only_answer: str = None - graph_vector_answer: str = None + raw_answer: Optional[str] = None + vector_only_answer: Optional[str] = None + graph_only_answer: Optional[str] = None + graph_vector_answer: Optional[str] = None - merged_result = None + merged_result: Optional[Any] = None def setup(self): self.schema = None @@ -184,9 +192,9 @@ def setup(self): self.graph_only_answer = None self.graph_vector_answer = None - self.vector_result = None - self.graph_result = None self.merged_result = None + + self.match_vids = None return CStatus() def to_json(self): diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 3f527f2fa..b881204f7 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -16,23 +16,21 @@ # under the License. -import json import os import traceback -from typing import Dict, Any, Union, Optional +from typing import Dict, Any, Union, List import gradio as gr from hugegraph_llm.flows.scheduler import SchedulerSingleton +from pyhugegraph.client import PyHugeClient from .embedding_utils import get_filename_prefix, get_index_folder_name -from .hugegraph_utils import get_hg_client, clean_hg_data +from .hugegraph_utils import clean_hg_data from .log import log from .vector_index_utils import read_documents from ..config import resource_path, huge_settings, llm_settings from ..indices.vector_index import VectorIndex from ..models.embeddings.init_embedding import Embeddings -from ..models.llms.init_llm import LLMs -from ..operators.kg_construction_task import KgBuilder def get_graph_index_info(): @@ -63,63 +61,29 @@ def clean_all_graph_index(): gr.Info("Clear graph index and text2gql index successfully!") -def clean_all_graph_data(): - clean_hg_data() - log.warning("Clear graph data successfully!") - gr.Info("Clear graph data successfully!") - - -def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: - schema = schema.strip() - if schema.startswith("{"): - try: - schema = json.loads(schema) - builder.import_schema(from_user_defined=schema) - except json.JSONDecodeError: - log.error("Invalid JSON format in schema. Please check it again.") - return "ERROR: Invalid JSON format in schema. Please check it carefully." +def get_vertex_details(vertex_ids: List[str], context: Dict) -> List[Dict[str, Any]]: + if isinstance(context.get("graph_client"), PyHugeClient): + client = context["graph_client"] else: - log.info("Get schema '%s' from graphdb.", schema) - builder.import_schema(from_hugegraph=schema) - return None + url = context.get("url") or "http://localhost:8080" + graph = context.get("graph") or "hugegraph" + user = context.get("user") or "admin" + pwd = context.get("pwd") or "admin" + gs = context.get("graphspace") or None + client = PyHugeClient(url, graph, user, pwd, gs) + if not vertex_ids: + return [] + formatted_ids = ", ".join(f"'{vid}'" for vid in vertex_ids) + gremlin_query = f"g.V({formatted_ids}).limit(20)" + result = client.gremlin().exec(gremlin=gremlin_query)["data"] + return result -def extract_graph_origin(input_file, input_text, schema, example_prompt) -> str: - texts = read_documents(input_file, input_text) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) - if not schema: - return "ERROR: please input with correct schema/format." - error_message = parse_schema(schema, builder) - if error_message: - return error_message - builder.chunk_split(texts, "document", "zh").extract_info( - example_prompt, "property_graph" - ) - - try: - context = builder.run() - if not context["vertices"] and not context["edges"]: - log.info("Please check the schema.(The schema may not match the Doc)") - return json.dumps( - { - "vertices": context["vertices"], - "edges": context["edges"], - "warning": "The schema may not match the Doc", - }, - ensure_ascii=False, - indent=2, - ) - return json.dumps( - {"vertices": context["vertices"], "edges": context["edges"]}, - ensure_ascii=False, - indent=2, - ) - except Exception as e: # pylint: disable=broad-exception-caught - log.error(e) - raise gr.Error(str(e)) +def clean_all_graph_data(): + clean_hg_data() + log.warning("Clear graph data successfully!") + gr.Info("Clear graph data successfully!") def extract_graph(input_file, input_text, schema, example_prompt) -> str: diff --git a/hugegraph-ml/pyproject.toml b/hugegraph-ml/pyproject.toml index 6d46ba74c..929eb3aa1 100644 --- a/hugegraph-ml/pyproject.toml +++ b/hugegraph-ml/pyproject.toml @@ -22,7 +22,7 @@ build-backend = "hatchling.build" [project] name = "hugegraph-ml" -version = "1.5.0" +version = "1.7.0" description = "Machine learning extensions for Apache HugeGraph." authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, diff --git a/hugegraph-python-client/pyproject.toml b/hugegraph-python-client/pyproject.toml index 81565d9ab..ddae125d8 100644 --- a/hugegraph-python-client/pyproject.toml +++ b/hugegraph-python-client/pyproject.toml @@ -17,7 +17,7 @@ [project] name = "hugegraph-python-client" -version = "1.5.0" +version = "1.7.0" description = "A Python SDK for Apache HugeGraph Database." authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, diff --git a/pyproject.toml b/pyproject.toml index 9e1624b49..7621b5fa6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ [project] name = "hugegraph-ai" -version = "1.5.0" +version = "1.7.0" description = "A repository for AI-related projects for Apache HugeGraph." authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, diff --git a/scripts/build_llm_image.sh b/scripts/build_llm_image.sh old mode 100644 new mode 100755 index 42aa36e39..7425b3df9 --- a/scripts/build_llm_image.sh +++ b/scripts/build_llm_image.sh @@ -18,7 +18,7 @@ set -e -tag="1.5.0" +tag="1.7.0" script_dir=$(realpath "$(dirname "$0")") diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/design.md b/spec/hugegraph-llm/fixed_flow/design.md similarity index 99% rename from .vibedev/spec/hugegraph-llm/fixed_flow/design.md rename to spec/hugegraph-llm/fixed_flow/design.md index c5777236d..2a9d08585 100644 --- a/.vibedev/spec/hugegraph-llm/fixed_flow/design.md +++ b/spec/hugegraph-llm/fixed_flow/design.md @@ -202,7 +202,7 @@ flowchart TD - `BuildVectorIndexFlow`: 向量索引构建工作流 - `GraphExtractFlow`: 图抽取工作流 - `ImportGraphDataFlow`: 图数据导入工作流 - - `UpdateVidEmbeddingsFlows`: 向量更新工作流 + - `UpdateVidEmbeddingsFlow`: 向量更新工作流 - `GetGraphIndexInfoFlow`: 图索引信息获取工作流 - `BuildSchemaFlow`: 模式构建工作流 - `PromptGenerateFlow`: 提示词生成工作流 diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/requirements.md b/spec/hugegraph-llm/fixed_flow/requirements.md similarity index 100% rename from .vibedev/spec/hugegraph-llm/fixed_flow/requirements.md rename to spec/hugegraph-llm/fixed_flow/requirements.md diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/tasks.md b/spec/hugegraph-llm/fixed_flow/tasks.md similarity index 100% rename from .vibedev/spec/hugegraph-llm/fixed_flow/tasks.md rename to spec/hugegraph-llm/fixed_flow/tasks.md diff --git a/vermeer-python-client/pyproject.toml b/vermeer-python-client/pyproject.toml index 986010899..d60acc075 100644 --- a/vermeer-python-client/pyproject.toml +++ b/vermeer-python-client/pyproject.toml @@ -17,7 +17,7 @@ [project] name = "vermeer-python-client" -version = "1.5.0" # Independently managed version for the vermeer-python-client package +version = "1.7.0" # Independently managed version for the vermeer-python-client package description = "A Python client library for interacting with Vermeer, a tool for managing and analyzing large-scale graph data." authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" } @@ -33,7 +33,7 @@ dependencies = [ "setuptools", "urllib3", "rich", - + # Vermeer specific dependencies "python-dateutil", ] From ff52f1191716438841f6e0daa4bf44e2b27ac140 Mon Sep 17 00:00:00 2001 From: jinglinwei Date: Tue, 14 Oct 2025 00:27:02 +0800 Subject: [PATCH 44/71] adopt some suggestions proposed by ai --- .../src/hugegraph_llm/flows/scheduler.py | 54 ++++++++++--------- .../nodes/hugegraph_node/fetch_graph_data.py | 11 ++-- .../nodes/hugegraph_node/graph_query_node.py | 18 ++++--- .../index_node/build_gremlin_example_index.py | 6 +-- .../hugegraph_llm/operators/operator_list.py | 2 + .../hugegraph_llm/utils/graph_index_utils.py | 7 ++- 6 files changed, 57 insertions(+), 41 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index 7cc5653d5..388c8c655 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -35,7 +35,7 @@ class Scheduler: - pipeline_pool: Dict[str, Any] = None + pipeline_pool: Dict[str, Any] max_pipeline: int def __init__(self, max_pipeline: int = 10): @@ -116,6 +116,7 @@ def schedule_flow(self, flow_name: str, *args, **kwargs): raise RuntimeError(error_msg) status = pipeline.run() if status.isErr(): + manager.add(pipeline) error_msg = f"Error in flow execution: {status.getInfo()}" log.error(error_msg) raise RuntimeError(error_msg) @@ -123,16 +124,18 @@ def schedule_flow(self, flow_name: str, *args, **kwargs): manager.add(pipeline) return res else: - # fetch pipeline & prepare input for flow - prepared_input = pipeline.getGParamWithNoEmpty("wkflow_input") - flow.prepare(prepared_input, *args, **kwargs) - status = pipeline.run() - if status.isErr(): - error_msg = f"Error in flow execution {status.getInfo()}" - log.error(error_msg) - raise RuntimeError(error_msg) - res = flow.post_deal(pipeline) - manager.release(pipeline) + try: + # fetch pipeline & prepare input for flow + prepared_input = pipeline.getGParamWithNoEmpty("wkflow_input") + flow.prepare(prepared_input, *args, **kwargs) + status = pipeline.run() + if status.isErr(): + error_msg = f"Error in flow execution {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + res = flow.post_deal(pipeline) + finally: + manager.release(pipeline) return res async def schedule_stream_flow(self, flow_name: str, *args, **kwargs): @@ -144,22 +147,21 @@ async def schedule_stream_flow(self, flow_name: str, *args, **kwargs): if pipeline is None: # call coresponding flow_func to create new workflow pipeline = flow.build_flow(*args, **kwargs) - try: - pipeline.getGParamWithNoEmpty("wkflow_input").stream = True - status = pipeline.init() - if status.isErr(): - error_msg = f"Error in flow init: {status.getInfo()}" - log.error(error_msg) - raise RuntimeError(error_msg) - status = pipeline.run() - if status.isErr(): - error_msg = f"Error in flow execution: {status.getInfo()}" - log.error(error_msg) - raise RuntimeError(error_msg) - async for res in flow.post_deal_stream(pipeline): - yield res - finally: + pipeline.getGParamWithNoEmpty("wkflow_input").stream = True + status = pipeline.init() + if status.isErr(): + error_msg = f"Error in flow init: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + status = pipeline.run() + if status.isErr(): manager.add(pipeline) + error_msg = f"Error in flow execution: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + async for res in flow.post_deal_stream(pipeline): + yield res + manager.add(pipeline) else: try: # fetch pipeline & prepare input for flow 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 078d7fdab..1e9e88a4e 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 @@ -13,6 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from PyCGraph import CStatus +from typing import Optional + +from hugegraph_llm.utils.log import log from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.hugegraph_op.fetch_graph_data import FetchGraphData from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -21,14 +25,15 @@ class FetchGraphDataNode(BaseNode): fetch_graph_data_op: FetchGraphData - context: WkFlowState = None - wk_input: WkFlowInput = None + context: Optional[WkFlowState] = None + wk_input: Optional[WkFlowInput] = None def node_init(self): try: client = get_hg_client() except Exception as e: - raise RuntimeError(f"can't initalzie HugeGraph client: {e}") + log.error("Failed to initialize HugeGraph client: %s", e) + return CStatus(-1, f"Failed to initialize HugeGraph client: {e}") self.fetch_graph_data_op = FetchGraphData(client) return super().node_init() diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py index 14f773db5..7bc9dab69 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py @@ -161,7 +161,7 @@ def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: ) except Exception as e: # pylint: disable=broad-except log.error(e) - context["graph_result"] = "" + context["graph_result"] = [] return context def _limit_property_query( @@ -405,9 +405,11 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: log.debug("Kneighbor gremlin query: %s", gremlin_query) paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( - self._format_graph_query_result(query_paths=paths) - ) + ( + graph_chain_knowledge, + vertex_degree_list, + knowledge_with_degree, + ) = self._format_graph_query_result(query_paths=paths) # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) if not graph_chain_knowledge: @@ -436,9 +438,11 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)[ "data" ] - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( - self._format_graph_query_result(query_paths=paths) - ) + ( + graph_chain_knowledge, + vertex_degree_list, + knowledge_with_degree, + ) = self._format_graph_query_result(query_paths=paths) context["graph_result"] = list(graph_chain_knowledge) if context["graph_result"]: diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py index 2237bfabb..8772959d7 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py @@ -30,10 +30,10 @@ class BuildGremlinExampleIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): - if self.wk_input.examples is not None: - examples = self.wk_input.examples - else: + if not self.wk_input.examples: return CStatus(-1, "examples is required in BuildGremlinExampleIndexNode") + examples = self.wk_input.examples + self.build_gremlin_example_index_op = BuildGremlinExampleIndex( get_embedding(llm_settings), examples ) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py index 46ddbd50e..129c36fc5 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py @@ -101,6 +101,8 @@ def print_result(self): return self def fetch_graph_data(self): + if self.graph is None: + raise ValueError("graph client is required for fetch_graph_data operation") self.operators.append(FetchGraphData(self.graph)) return self diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index b881204f7..a0e376ac6 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -61,7 +61,9 @@ def clean_all_graph_index(): gr.Info("Clear graph index and text2gql index successfully!") -def get_vertex_details(vertex_ids: List[str], context: Dict) -> List[Dict[str, Any]]: +def get_vertex_details( + vertex_ids: List[str], context: Dict[str, Any] +) -> List[Dict[str, Any]]: if isinstance(context.get("graph_client"), PyHugeClient): client = context["graph_client"] else: @@ -128,5 +130,6 @@ def build_schema(input_text, query_example, few_shot): return scheduler.schedule_flow( "build_schema", input_text, query_example, few_shot ) - except (TypeError, ValueError) as e: + except Exception as e: # pylint: disable=broad-exception-caught + log.error("Schema generation failed: %s", e) raise gr.Error(f"Schema generation failed: {e}") From 1d56c5ce21dfa09ff6588fff112201f10e4e62bd Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 15:59:33 +0800 Subject: [PATCH 45/71] Merge upstream/agenticrag/dev: integrate flows/nodes architecture + pluggable vector stores - Adopt new GPipeline/Scheduler architecture from agenticrag/dev - Preserve pluggable vector index (Faiss/Milvus/Qdrant) from vector-db - Keep Embeddings() unified interface for embedding models - Remove deprecated operators: gremlin_generate_task, graph_rag_query, kg_construction_task - Align log formatting (%s instead of f-string) - Return dict instead of json.dumps in flows for consistency --- hugegraph-llm/README.md | 3 + hugegraph-llm/config.md | 13 + hugegraph-llm/pyproject.toml | 32 +- .../src/hugegraph_llm/api/admin_api.py | 5 +- .../src/hugegraph_llm/config/__init__.py | 8 +- .../src/hugegraph_llm/config/generate.py | 9 +- .../hugegraph_llm/config/hugegraph_config.py | 23 +- .../src/hugegraph_llm/config/index_config.py | 38 ++ .../src/hugegraph_llm/config/llm_config.py | 77 ++- .../hugegraph_llm/config/models/__init__.py | 2 + .../demo/rag_demo/admin_block.py | 3 +- .../src/hugegraph_llm/demo/rag_demo/app.py | 3 +- .../demo/rag_demo/configs_block.py | 146 +++++- .../hugegraph_llm/demo/rag_demo/rag_block.py | 196 ++++---- .../demo/rag_demo/text2gremlin_block.py | 210 +++++--- .../demo/rag_demo/vector_graph_block.py | 50 +- .../src/hugegraph_llm/flows/common.py | 8 +- .../flows/get_graph_index_info.py | 34 +- .../flows/rag_flow_graph_only.py | 10 +- .../flows/rag_flow_graph_vector.py | 26 +- .../src/hugegraph_llm/flows/rag_flow_raw.py | 6 +- .../flows/rag_flow_vector_only.py | 10 +- .../src/hugegraph_llm/flows/scheduler.py | 4 +- .../src/hugegraph_llm/indices/graph_index.py | 1 + .../src/hugegraph_llm/indices/vector_index.py | 170 ------- .../indices/vector_index/base.py | 107 ++++ .../vector_index/faiss_vector_store.py | 140 ++++++ .../vector_index/milvus_vector_store.py | 263 ++++++++++ .../vector_index/qdrant_vector_store.py | 221 ++++++++ .../hugegraph_llm/middleware/middleware.py | 2 +- .../hugegraph_llm/models/embeddings/base.py | 26 +- .../models/embeddings/init_embedding.py | 39 +- .../models/embeddings/litellm.py | 43 +- .../hugegraph_llm/models/embeddings/ollama.py | 60 ++- .../hugegraph_llm/models/embeddings/openai.py | 22 +- .../nodes/common_node/merge_rerank_node.py | 16 +- .../nodes/hugegraph_node/graph_query_node.py | 471 ++---------------- .../nodes/index_node/build_semantic_index.py | 5 +- .../nodes/index_node/build_vector_index.py | 5 +- .../index_node/gremlin_example_index_query.py | 9 +- .../index_node/semantic_id_query_node.py | 35 +- .../nodes/index_node/vector_query_node.py | 13 +- .../nodes/llm_node/answer_synthesize_node.py | 8 +- .../nodes/llm_node/keyword_extract_node.py | 16 +- .../nodes/llm_node/text2gremlin.py | 13 +- .../operators/common_op/check_schema.py | 2 +- .../hugegraph_op/fetch_graph_data.py | 1 - .../operators/hugegraph_op/schema_manager.py | 2 +- .../index_op/build_gremlin_example_index.py | 32 +- .../index_op/build_semantic_index.py | 67 +-- .../operators/index_op/build_vector_index.py | 34 +- .../index_op/gremlin_example_index_query.py | 69 ++- .../operators/index_op/semantic_id_query.py | 25 +- .../operators/index_op/vector_index_query.py | 21 +- .../operators/llm_op/answer_synthesize.py | 54 +- .../operators/llm_op/disambiguate_data.py | 3 +- .../operators/llm_op/gremlin_generate.py | 7 +- .../operators/llm_op/info_extract.py | 11 +- .../llm_op/property_graph_extract.py | 3 +- .../hugegraph_llm/operators/operator_list.py | 2 +- .../src/hugegraph_llm/state/ai_state.py | 12 +- .../src/hugegraph_llm/utils/anchor.py | 3 +- hugegraph-llm/src/hugegraph_llm/utils/log.py | 0 .../hugegraph_llm/utils/vector_index_utils.py | 86 ++-- ...or_index.py => test_faiss_vector_index.py} | 10 +- .../tests/indices/test_milvus_vector_index.py | 100 ++++ .../tests/indices/test_qdrant_vector_index.py | 102 ++++ hugegraph-ml/src/hugegraph_ml/models/bgrl.py | 2 +- .../src/pyhugegraph/api/auth.py | 12 +- .../src/pyhugegraph/api/graph.py | 8 +- .../src/pyhugegraph/api/gremlin.py | 4 +- .../src/pyhugegraph/api/schema.py | 8 +- .../src/pyhugegraph/api/traverser.py | 8 +- .../src/pyhugegraph/utils/util.py | 7 +- 74 files changed, 1913 insertions(+), 1383 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/config/index_config.py delete mode 100644 hugegraph-llm/src/hugegraph_llm/indices/vector_index.py create mode 100644 hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py create mode 100644 hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py create mode 100644 hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py create mode 100644 hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py mode change 100755 => 100644 hugegraph-llm/src/hugegraph_llm/utils/log.py rename hugegraph-llm/src/tests/indices/{test_vector_index.py => test_faiss_vector_index.py} (81%) create mode 100644 hugegraph-llm/src/tests/indices/test_milvus_vector_index.py create mode 100644 hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py diff --git a/hugegraph-llm/README.md b/hugegraph-llm/README.md index f0eeb3136..c91e593e2 100644 --- a/hugegraph-llm/README.md +++ b/hugegraph-llm/README.md @@ -113,6 +113,9 @@ python -m hugegraph_llm.demo.rag_demo.app --host 127.0.0.1 --port 18001 > The following commands assume you're in the activated virtual environment from step 4 above ```bash +# To use vector database backends (e.g., Milvus, Qdrant), sync the optional dependencies: +uv sync --extra vectordb + # Download NLTK stopwords for better text processing python ./src/hugegraph_llm/operators/common_op/nltk_helper.py diff --git a/hugegraph-llm/config.md b/hugegraph-llm/config.md index a55172f33..5b0e766d5 100644 --- a/hugegraph-llm/config.md +++ b/hugegraph-llm/config.md @@ -16,6 +16,7 @@ - [LiteLLM 配置](#litellm-配置) - [重排序配置](#重排序配置) - [HugeGraph 数据库配置](#hugegraph-数据库配置) + - [向量数据库配置](#向量数据库配置) - [管理员配置](#管理员配置) - [配置使用示例](#配置使用示例) - [配置文件位置](#配置文件位置) @@ -127,6 +128,18 @@ | `TOPK_PER_KEYWORD` | Optional[Integer] | 1 | 每个关键词返回的 TopK 数量 | | `TOPK_RETURN_RESULTS` | Optional[Integer] | 20 | 返回结果数量 | +### 向量数据库配置 + +| 配置项 | 类型 | 默认值 | 说明 | +|------------------|------------------|-------|------------------------| +| `QDRANT_HOST` | Optional[String] | None | Qdrant 服务器主机地址 | +| `QDRANT_PORT` | Integer | 6333 | Qdrant 服务器端口 | +| `QDRANT_API_KEY` | Optional[String] | None | Qdrant API 密钥（如果设置了的话） | +| `MILVUS_HOST` | Optional[String] | None | Milvus 服务器主机地址 | +| `MILVUS_PORT` | Integer | 19530 | Milvus 服务器端口 | +| `MILVUS_USER` | String | "" | Milvus 用户名 | +| `MILVUS_PASSWORD`| String | "" | Milvus 密码 | + ### 管理员配置 | 配置项 | 类型 | 默认值 | 说明 | diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index f5301591b..a46836754 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -22,11 +22,12 @@ description = "A tool for the implementation and research related to large langu authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, ] +maintainers = [ + { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, +] readme = "README.md" license = "Apache-2.0" requires-python = ">=3.10,<3.12" - - dependencies = [ # Common dependencies "decorator", @@ -39,6 +40,7 @@ dependencies = [ "numpy", "pandas", "pydantic", + "tqdm", # LLM specific dependencies "openai", @@ -60,6 +62,13 @@ dependencies = [ "hugegraph-python-client", "pycgraph", ] + +[project.optional-dependencies] +vectordb = [ + "pymilvus==2.5.9", + "qdrant-client==1.14.2", +] + [project.urls] homepage = "https://hugegraph.apache.org/" repository = "https://github.com/apache/incubator-hugegraph-ai" @@ -87,3 +96,22 @@ allow-direct-references = true [tool.uv.sources] hugegraph-python-client = { workspace = true } pycgraph = { git = "https://github.com/ChunelFeng/CGraph.git", subdirectory = "python", rev = "main", marker = "sys_platform == 'linux'" } + +[tool.mypy] +disable_error_code = ["import-untyped"] +check_untyped_defs = true +disallow_untyped_defs = false + +[tool.ruff] +line-length = 120 +indent-width = 4 +extend-exclude = [] + +# TODO: move this config in the root pyproject.toml & add more rules for it +[tool.ruff.lint] +extend-select = ["I"] + +[tool.ruff.format] +quote-style = "preserve" +indent-style = "space" +line-ending = "auto" diff --git a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py index 4c192c29c..788c62082 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py @@ -33,9 +33,8 @@ async def log_stream_api(req: LogStreamRequest): if admin_settings.admin_token != req.admin_token: raise generate_response( RAGResponse( - status_code=status.HTTP_403_FORBIDDEN, # pylint: disable=E0702 - message="Invalid admin_token", - ) + status_code=status.HTTP_403_FORBIDDEN, message="Invalid admin_token" + ) # pylint: disable=E0702 ) log_path = os.path.join("logs", req.log_file) diff --git a/hugegraph-llm/src/hugegraph_llm/config/__init__.py b/hugegraph-llm/src/hugegraph_llm/config/__init__.py index 5d4f5782d..43efb0ab3 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/__init__.py +++ b/hugegraph-llm/src/hugegraph_llm/config/__init__.py @@ -16,14 +16,15 @@ # under the License. -__all__ = ["huge_settings", "admin_settings", "llm_settings", "resource_path"] +__all__ = ["huge_settings", "admin_settings", "llm_settings", "resource_path", "index_settings"] import os -from .prompt_config import PromptConfig -from .hugegraph_config import HugeGraphConfig from .admin_config import AdminConfig +from .hugegraph_config import HugeGraphConfig +from .index_config import IndexConfig from .llm_config import LLMConfig +from .prompt_config import PromptConfig llm_settings = LLMConfig() prompt = PromptConfig(llm_settings) @@ -31,6 +32,7 @@ huge_settings = HugeGraphConfig() admin_settings = AdminConfig() +index_settings = IndexConfig() package_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) resource_path = os.path.join(package_path, "resources") diff --git a/hugegraph-llm/src/hugegraph_llm/config/generate.py b/hugegraph-llm/src/hugegraph_llm/config/generate.py index 4b40e899f..1bd7adea8 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/generate.py +++ b/hugegraph-llm/src/hugegraph_llm/config/generate.py @@ -18,7 +18,13 @@ import argparse -from hugegraph_llm.config import huge_settings, admin_settings, llm_settings, PromptConfig +from hugegraph_llm.config import ( + PromptConfig, + admin_settings, + huge_settings, + index_settings, + llm_settings, +) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Generate hugegraph-llm config file") @@ -30,4 +36,5 @@ huge_settings.generate_env() admin_settings.generate_env() llm_settings.generate_env() + index_settings.generate_env() PromptConfig(llm_settings).generate_yaml_file() diff --git a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py b/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py index 69abf0fbc..1937eda10 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py @@ -16,6 +16,7 @@ # under the License. from typing import Optional + from .models import BaseConfig @@ -23,21 +24,21 @@ class HugeGraphConfig(BaseConfig): """HugeGraph settings""" # graph server config - graph_url: Optional[str] = "127.0.0.1:8080" - graph_name: Optional[str] = "hugegraph" - graph_user: Optional[str] = "admin" - graph_pwd: Optional[str] = "xxx" + graph_url: str = "127.0.0.1:8080" + graph_name: str = "hugegraph" + graph_user: str = "admin" + graph_pwd: str = "xxx" graph_space: Optional[str] = None # graph query config - limit_property: Optional[str] = "False" - max_graph_path: Optional[int] = 10 - max_graph_items: Optional[int] = 30 - edge_limit_pre_label: Optional[int] = 8 + limit_property: str = "False" + max_graph_path: int = 10 + max_graph_items: int = 30 + edge_limit_pre_label: int = 8 # vector config - vector_dis_threshold: Optional[float] = 0.9 - topk_per_keyword: Optional[int] = 1 + vector_dis_threshold: float = 0.9 + topk_per_keyword: int = 1 # rerank config - topk_return_results: Optional[int] = 20 + topk_return_results: int = 20 diff --git a/hugegraph-llm/src/hugegraph_llm/config/index_config.py b/hugegraph-llm/src/hugegraph_llm/config/index_config.py new file mode 100644 index 000000000..63895e6a7 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/config/index_config.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. + +import os +from typing import Optional + +from .models import BaseConfig + + +class IndexConfig(BaseConfig): + """LLM settings""" + + qdrant_host: Optional[str] = os.environ.get("QDRANT_HOST", None) + qdrant_port: int = int(os.environ.get("QDRANT_PORT", "6333")) + qdrant_api_key: Optional[str] = ( + os.environ.get("QDRANT_API_KEY") if os.environ.get("QDRANT_API_KEY") else None + ) + + milvus_host: Optional[str] = os.environ.get("MILVUS_HOST", None) + milvus_port: int = int(os.environ.get("MILVUS_PORT", "19530")) + milvus_user: str = os.environ.get("MILVUS_USER", "") + milvus_password: str = os.environ.get("MILVUS_PASSWORD", "") + + cur_vector_index: str = os.environ.get("CUR_VECTOR_INDEX", "Faiss") diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index 916d70ddf..75494ab31 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -17,7 +17,7 @@ import os -from typing import Optional, Literal +from typing import Literal, Optional from .models import BaseConfig @@ -32,61 +32,60 @@ class LLMConfig(BaseConfig): embedding_type: Optional[Literal["openai", "litellm", "ollama/local"]] = "openai" reranker_type: Optional[Literal["cohere", "siliconflow"]] = None # 1. OpenAI settings - openai_chat_api_base: Optional[str] = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) - openai_chat_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") - openai_chat_language_model: Optional[str] = "gpt-4.1-mini" - openai_extract_api_base: Optional[str] = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) - openai_extract_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") - openai_extract_language_model: Optional[str] = "gpt-4.1-mini" - openai_text2gql_api_base: Optional[str] = os.environ.get( - "OPENAI_BASE_URL", "https://api.openai.com/v1" - ) - openai_text2gql_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") - openai_text2gql_language_model: Optional[str] = "gpt-4.1-mini" - openai_embedding_api_base: Optional[str] = os.environ.get( + openai_chat_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_chat_api_key: str | None = os.environ.get("OPENAI_API_KEY") + openai_chat_language_model: str = "gpt-4.1-mini" + openai_extract_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_extract_api_key: str | None = os.environ.get("OPENAI_API_KEY") + openai_extract_language_model: str = "gpt-4.1-mini" + openai_text2gql_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_text2gql_api_key: str | None = os.environ.get("OPENAI_API_KEY") + openai_text2gql_language_model: str = "gpt-4.1-mini" + openai_embedding_api_base: str = os.environ.get( "OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1" ) - openai_embedding_api_key: Optional[str] = os.environ.get("OPENAI_EMBEDDING_API_KEY") - openai_embedding_model: Optional[str] = "text-embedding-3-small" + openai_embedding_api_key: str | None = os.environ.get("OPENAI_EMBEDDING_API_KEY") + openai_embedding_model: str = "text-embedding-3-small" + openai_embedding_model_dim: int = 1536 openai_chat_tokens: int = 8192 openai_extract_tokens: int = 256 openai_text2gql_tokens: int = 4096 # 2. Rerank settings - cohere_base_url: Optional[str] = os.environ.get( - "CO_API_URL", "https://api.cohere.com/v1/rerank" - ) + cohere_base_url: str = os.environ.get("CO_API_URL", "https://api.cohere.com/v1/rerank") reranker_api_key: Optional[str] = None reranker_model: Optional[str] = None # 3. Ollama settings - ollama_chat_host: Optional[str] = "127.0.0.1" - ollama_chat_port: Optional[int] = 11434 - ollama_chat_language_model: Optional[str] = None - ollama_extract_host: Optional[str] = "127.0.0.1" - ollama_extract_port: Optional[int] = 11434 - ollama_extract_language_model: Optional[str] = None - ollama_text2gql_host: Optional[str] = "127.0.0.1" - ollama_text2gql_port: Optional[int] = 11434 - ollama_text2gql_language_model: Optional[str] = None - ollama_embedding_host: Optional[str] = "127.0.0.1" - ollama_embedding_port: Optional[int] = 11434 - ollama_embedding_model: Optional[str] = None - # 4. LiteLLM settings + ollama_chat_host: str = "127.0.0.1" + ollama_chat_port: int = 11434 + ollama_chat_language_model: str | None = None + ollama_extract_host: str = "127.0.0.1" + ollama_extract_port: int = 11434 + ollama_extract_language_model: str | None = None + ollama_text2gql_host: str = "127.0.0.1" + ollama_text2gql_port: int = 11434 + ollama_text2gql_language_model: str | None = None + ollama_embedding_host: str = "127.0.0.1" + ollama_embedding_port: int = 11434 + ollama_embedding_model: str = "quentinz/bge-large-zh-v1.5" + _env_ollama_dim = os.getenv("OLLAMA_EMBEDDING_MODEL_DIM") + ollama_embedding_model_dim: Optional[int] = int(_env_ollama_dim) if _env_ollama_dim else None + + # 4. QianFan/WenXin settings (removed) + + # 5. LiteLLM settings litellm_chat_api_key: Optional[str] = None litellm_chat_api_base: Optional[str] = None - litellm_chat_language_model: Optional[str] = "openai/gpt-4.1-mini" + litellm_chat_language_model: str = "openai/gpt-4.1-mini" litellm_chat_tokens: int = 8192 litellm_extract_api_key: Optional[str] = None litellm_extract_api_base: Optional[str] = None - litellm_extract_language_model: Optional[str] = "openai/gpt-4.1-mini" + litellm_extract_language_model: str = "openai/gpt-4.1-mini" litellm_extract_tokens: int = 256 litellm_text2gql_api_key: Optional[str] = None litellm_text2gql_api_base: Optional[str] = None - litellm_text2gql_language_model: Optional[str] = "openai/gpt-4.1-mini" + litellm_text2gql_language_model: str = "openai/gpt-4.1-mini" litellm_text2gql_tokens: int = 4096 litellm_embedding_api_key: Optional[str] = None litellm_embedding_api_base: Optional[str] = None - litellm_embedding_model: Optional[str] = "openai/text-embedding-3-small" + litellm_embedding_model: str = "openai/text-embedding-3-small" + litellm_embedding_model_dim: int = 1536 diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py b/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py index e73646fd1..087d89477 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/__init__.py @@ -17,3 +17,5 @@ from .base_config import BaseConfig from .base_prompt_config import BasePromptConfig + +__all__ = ["BaseConfig", "BasePromptConfig"] diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 1b2032b23..0b48a84dc 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -19,7 +19,6 @@ from collections import deque import gradio as gr -from gradio import Request from hugegraph_llm.config import admin_settings from hugegraph_llm.utils.log import log @@ -70,7 +69,7 @@ def clear_llm_server_log(): # Function to validate password and control access to logs -def check_password(password, request: Request = None): +def check_password(password, request=None): client_ip = request.client.host if request else "Unknown IP" admin_token = admin_settings.admin_token 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 eabd5dd9b..b0979763c 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -115,7 +115,7 @@ def init_rag_ui() -> gr.Interface: def refresh_ui_config_prompt() -> tuple: # we can use its __init__() for in-place reload # settings.from_env() - huge_settings.__init__() # pylint: disable=C2801 + huge_settings.__init__() # type: ignore[misc] # pylint: disable=C2801 prompt.ensure_yaml_file_exists() return ( huge_settings.graph_url, @@ -164,6 +164,7 @@ def create_app(): # we don't need to manually check the env now # settings.check_env() prompt.update_yaml_file() + assert admin_settings.enable_login auth_enabled = admin_settings.enable_login.lower() == "true" log.info("(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled") api_auth = APIRouter(dependencies=[Depends(authenticate)] if auth_enabled else []) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 8c595c30d..a6f57fc67 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -25,7 +25,7 @@ from dotenv import dotenv_values from requests.auth import HTTPBasicAuth -from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.config import huge_settings, index_settings, llm_settings from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding from hugegraph_llm.models.llms.litellm import LiteLLMClient from hugegraph_llm.utils.log import log @@ -106,6 +106,83 @@ def test_api_connection( return resp.status_code +def apply_vector_engine(engine: str): + # Persist the vector engine selection + setattr(index_settings, "cur_vector_index", engine) + try: + index_settings.update_env() + except Exception: # pylint: disable=W0718 + pass + gr.Info("Configured!") + + +def apply_vector_engine_backend( + engine: str, + host: Optional[str] = None, + port: Optional[str] = None, + user: Optional[str] = None, + password: Optional[str] = None, + api_key: Optional[str] = None, + origin_call=None, +) -> int: + """Test connection and persist per-engine connection settings""" + status_code = -1 + + # Test connection first + try: + if engine == "Milvus": + from pymilvus import connections, utility + + connections.connect( + host=host, port=int(port or 19530), user=user or "", password=password or "" + ) + # Test if we can list collections + _ = utility.list_collections() + connections.disconnect("default") + status_code = 200 + elif engine == "Qdrant": + from qdrant_client import QdrantClient + + client = QdrantClient(host=host, port=int(port or 6333), api_key=api_key) + # Test if we can get collections + _ = client.get_collections() + status_code = 200 + except ImportError as e: + msg = f"Missing dependency: {e}. Please install with: uv sync --extra vectordb" + if origin_call is None: + raise gr.Error(msg) from e + return -1 + except Exception as e: + msg = f"Connection failed: {e}" + log.error(msg) + if origin_call is None: + raise gr.Error(msg) from e + return -1 + + # Persist settings after successful test + if engine == "Milvus": + if host is not None: + index_settings.milvus_host = host + if port is not None and str(port).strip(): + index_settings.milvus_port = int(port) # type: ignore[arg-type] + index_settings.milvus_user = user or "" + index_settings.milvus_password = password or "" + elif engine == "Qdrant": + if host is not None: + index_settings.qdrant_host = host + if port is not None and str(port).strip(): + index_settings.qdrant_port = int(port) # type: ignore[arg-type] + # Empty string treated as None for api key + index_settings.qdrant_api_key = api_key or None + + try: + index_settings.update_env() + except Exception: # pylint: disable=W0718 + pass + gr.Info("Configured!") + return status_code + + def apply_embedding_config(arg1, arg2, arg3, origin_call=None) -> int: status_code = -1 embedding_option = llm_settings.embedding_type @@ -561,8 +638,14 @@ def embedding_settings(embedding_type): elif embedding_type == "ollama/local": with gr.Row(): embedding_config_input = [ - gr.Textbox(value=llm_settings.ollama_embedding_host, label="host"), - gr.Textbox(value=str(llm_settings.ollama_embedding_port), label="port"), + gr.Textbox( + value=llm_settings.ollama_embedding_host, + label="host", + ), + gr.Textbox( + value=str(llm_settings.ollama_embedding_port), + label="port", + ), gr.Textbox( value=llm_settings.ollama_embedding_model, label="model_name", @@ -596,7 +679,6 @@ def embedding_settings(embedding_type): embedding_config_button = gr.Button("Apply Configuration") - # Call the separate apply_embedding_configuration function here embedding_config_button.click( # pylint: disable=no-member fn=apply_embedding_config, inputs=embedding_config_input, # pylint: disable=no-member @@ -616,18 +698,18 @@ def reranker_settings(reranker_type): with gr.Row(): reranker_config_input = [ gr.Textbox( - value=llm_settings.reranker_api_key, + value=lambda: llm_settings.reranker_api_key, label="api_key", type="password", ), - gr.Textbox(value=llm_settings.reranker_model, label="model"), - gr.Textbox(value=llm_settings.cohere_base_url, label="base_url"), + gr.Textbox(value=lambda: llm_settings.reranker_model, label="model"), + gr.Textbox(value=lambda: llm_settings.cohere_base_url, label="base_url"), ] elif reranker_type == "siliconflow": with gr.Row(): reranker_config_input = [ gr.Textbox( - value=llm_settings.reranker_api_key, + value=lambda: llm_settings.reranker_api_key, label="api_key", type="password", ), @@ -653,6 +735,54 @@ def reranker_settings(reranker_type): inputs=reranker_config_input, # pylint: disable=no-member ) + with gr.Accordion("5. Set up the vector engine.", open=False): + engine_selector = gr.Dropdown( + choices=["Faiss", "Milvus", "Qdrant"], + value=index_settings.cur_vector_index, + label="Select vector engine.", + ) + engine_selector.select( + fn=lambda engine: setattr(index_settings, "cur_vector_index", engine), + inputs=[engine_selector], + ) + + @gr.render(inputs=[engine_selector]) + def vector_engine_settings(engine): + if engine == "Milvus": + with gr.Row(): + milvus_inputs = [ + gr.Textbox(value=index_settings.milvus_host, label="host"), + gr.Textbox(value=str(index_settings.milvus_port), label="port"), + gr.Textbox(value=index_settings.milvus_user, label="user"), + gr.Textbox( + value=index_settings.milvus_password, label="password", type="password" + ), + ] + apply_backend_button = gr.Button("Apply Configuration") + apply_backend_button.click( + partial(apply_vector_engine_backend, "Milvus"), inputs=milvus_inputs + ) + elif engine == "Qdrant": + with gr.Row(): + qdrant_inputs = [ + gr.Textbox(value=index_settings.qdrant_host, label="host"), + gr.Textbox(value=str(index_settings.qdrant_port), label="port"), + gr.Textbox( + value=(index_settings.qdrant_api_key or ""), + label="api_key", + type="password", + ), + ] + apply_backend_button = gr.Button("Apply Configuration") + apply_backend_button.click( + lambda h, p, k: apply_vector_engine_backend("Qdrant", h, p, None, None, k), + inputs=qdrant_inputs, + ) + else: + gr.Markdown("✅ Faiss 本地索引无需额外配置。") + apply_faiss_button = gr.Button("Apply Configuration") + apply_faiss_button.click(lambda: apply_vector_engine(engine)) + # The reason for returning this partial value is the functional need to refresh the ui return graph_config_input diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 60ca6ae55..5482cdd87 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -18,14 +18,20 @@ # pylint: disable=E1101 import os -from typing import AsyncGenerator, Tuple, Literal, Optional +from typing import Any, AsyncGenerator, Literal, Optional, Tuple import gradio as gr -from hugegraph_llm.flows.scheduler import SchedulerSingleton import pandas as pd -from gradio.utils import NamedString -from hugegraph_llm.config import resource_path, prompt, llm_settings +from hugegraph_llm.config import ( + huge_settings, + index_settings, + llm_settings, + prompt, + resource_path, +) +from hugegraph_llm.operators.graph_rag_task import RAGPipeline +from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize from hugegraph_llm.utils.decorators import with_task_id from hugegraph_llm.utils.log import log @@ -67,51 +73,49 @@ def rag_answer( gr.Warning("Please select at least one generate mode.") return "", "", "", "" - scheduler = SchedulerSingleton.get_instance() - try: - # Select workflow by mode to avoid fetching the wrong pipeline from the pool - if graph_vector_answer or (graph_only_answer and vector_only_answer): - flow_key = "rag_graph_vector" - elif vector_only_answer: - flow_key = "rag_vector_only" - elif graph_only_answer: - flow_key = "rag_graph_only" - elif raw_answer: - flow_key = "rag_raw" - else: - raise RuntimeError("Unsupported flow type") + rag = RAGPipeline() + if vector_search: + rag.query_vector_index(vector_index_str=index_settings.cur_vector_index) + if graph_search: + rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid( + vector_index_str=index_settings.cur_vector_index, + vector_dis_threshold=vector_dis_threshold, + topk_per_keyword=topk_per_keyword, + ).import_schema(huge_settings.graph_name).query_graphdb( + num_gremlin_generate_example=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, + max_graph_items=max_graph_items, + ) + # TODO: add more user-defined search strategies + rag.merge_dedup_rerank( + graph_ratio=graph_ratio, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + topk_return_results=topk_return_results, + ) + rag.synthesize_answer( + raw_answer, + vector_only_answer, + graph_only_answer, + graph_vector_answer, + answer_prompt, + ) - res = scheduler.schedule_flow( - flow_key, + try: + context = rag.run( + verbose=True, query=text, vector_search=vector_search, graph_search=graph_search, - raw_answer=raw_answer, - vector_only_answer=vector_only_answer, - graph_only_answer=graph_only_answer, - graph_vector_answer=graph_vector_answer, - graph_ratio=graph_ratio, - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - custom_related_information=custom_related_information, - answer_prompt=answer_prompt, - keywords_extract_prompt=keywords_extract_prompt, - gremlin_tmpl_num=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, max_graph_items=max_graph_items, - topk_return_results=topk_return_results, - vector_dis_threshold=vector_dis_threshold, - topk_per_keyword=topk_per_keyword, ) - if res.get("switch_to_bleu"): - gr.Warning( - "Online reranker fails, automatically switches to local bleu rerank." - ) + if context.get("switch_to_bleu"): + gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") return ( - res.get("raw_answer", ""), - res.get("vector_only_answer", ""), - res.get("graph_only_answer", ""), - res.get("graph_vector_answer", ""), + context.get("raw_answer", ""), + context.get("vector_only_answer", ""), + context.get("graph_only_answer", ""), + context.get("graph_vector_answer", ""), ) except ValueError as e: log.critical(e) @@ -185,47 +189,47 @@ async def rag_answer_streaming( yield "", "", "", "" return - try: - # Select the specific streaming workflow - scheduler = SchedulerSingleton.get_instance() - if graph_vector_answer or (graph_only_answer and vector_only_answer): - flow_key = "rag_graph_vector" - elif vector_only_answer: - flow_key = "rag_vector_only" - elif graph_only_answer: - flow_key = "rag_graph_only" - elif raw_answer: - flow_key = "rag_raw" - else: - raise RuntimeError("Unsupported flow type") + rag = RAGPipeline() + if vector_search: + rag.query_vector_index(vector_index_str=index_settings.cur_vector_index) + if graph_search: + rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid( + vector_index_str=index_settings.cur_vector_index + ).import_schema(huge_settings.graph_name).query_graphdb( + num_gremlin_generate_example=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, + ) + rag.merge_dedup_rerank( + graph_ratio, + rerank_method, + near_neighbor_first, + ) + # rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) - async for res in scheduler.schedule_stream_flow( - flow_key, + try: + context = rag.run( + verbose=True, query=text, vector_search=vector_search, graph_search=graph_search, + ) + if context.get("switch_to_bleu"): + gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") + answer_synthesize = AnswerSynthesize( raw_answer=raw_answer, vector_only_answer=vector_only_answer, graph_only_answer=graph_only_answer, graph_vector_answer=graph_vector_answer, - graph_ratio=graph_ratio, - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - custom_related_information=custom_related_information, - answer_prompt=answer_prompt, - keywords_extract_prompt=keywords_extract_prompt, - gremlin_tmpl_num=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - ): - if res.get("switch_to_bleu"): - gr.Warning( - "Online reranker fails, automatically switches to local bleu rerank." - ) + prompt_template=answer_prompt, + ) + async for context in answer_synthesize.run_streaming(context): + if context.get("switch_to_bleu"): + gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") yield ( - res.get("raw_answer", ""), - res.get("vector_only_answer", ""), - res.get("graph_only_answer", ""), - res.get("graph_vector_answer", ""), + context.get("raw_answer", ""), + context.get("vector_only_answer", ""), + context.get("graph_only_answer", ""), + context.get("graph_vector_answer", ""), ) except ValueError as e: log.critical(e) @@ -290,9 +294,7 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): - raw_radio = gr.Radio( - choices=[True, False], value=False, label="Basic LLM Answer" - ) + raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") vector_only_radio = gr.Radio( choices=[True, False], value=False, label="Vector-only Answer" ) @@ -363,7 +365,7 @@ def toggle_slider(enable): """## 2. (Batch) Back-testing ) > 1. Download the template file & fill in the questions you want to test. > 2. Upload the file & click the button to generate answers. (Preview shows the first 40 lines) - > 3. The answer options are the same as the above RAG/Q&A frame + > 3. The answer options are the same as the above RAG/Q&A frame """ ) tests_df_headers = [ @@ -377,11 +379,9 @@ def toggle_slider(enable): # FIXME: "demo" might conflict with the graph name, it should be modified. answers_path = os.path.join(resource_path, "demo", "questions_answers.xlsx") questions_path = os.path.join(resource_path, "demo", "questions.xlsx") - questions_template_path = os.path.join( - resource_path, "demo", "questions_template.xlsx" - ) + questions_template_path = os.path.join(resource_path, "demo", "questions_template.xlsx") - def read_file_to_excel(file: NamedString, line_count: Optional[int] = None): + def read_file_to_excel(file: Any, line_count: Optional[int] = None): df = None if not file: return pd.DataFrame(), 1 @@ -389,15 +389,15 @@ def read_file_to_excel(file: NamedString, line_count: Optional[int] = None): df = pd.read_excel(file.name, nrows=line_count) if file else pd.DataFrame() elif file.name.endswith(".csv"): df = pd.read_csv(file.name, nrows=line_count) if file else pd.DataFrame() - df.to_excel(questions_path, index=False) - if df.empty: + df.to_excel(questions_path, index=False) # type:ignore + if df.empty: # type:ignore df = pd.DataFrame([[""] * len(tests_df_headers)], columns=tests_df_headers) else: - df.columns = tests_df_headers + df.columns = tests_df_headers # type:ignore # truncate the dataframe if it's too long - if len(df) > 40: - return df.head(40), 40 - return df, len(df) + if len(df) > 40: # type:ignore + return df.head(40), 40 # type:ignore + return df, len(df) # type:ignore def change_showing_excel(line_count): if os.path.exists(answers_path): @@ -459,18 +459,12 @@ def several_rag_answer( file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" ) with gr.Column(): - test_template_file = os.path.join( - resource_path, "demo", "questions_template.xlsx" - ) + test_template_file = os.path.join(resource_path, "demo", "questions_template.xlsx") gr.File(value=test_template_file, label="Download Template File") - answer_max_line_count = gr.Number( - 1, label="Max Lines To Show", minimum=1, maximum=40 - ) + answer_max_line_count = gr.Number(1, label="Max Lines To Show", minimum=1, maximum=40) answers_btn = gr.Button("Generate Answer (Batch)", variant="primary") # TODO: Set individual progress bars for dataframe - qa_dataframe = gr.DataFrame( - label="Questions & Answers (Preview)", headers=tests_df_headers - ) + qa_dataframe = gr.DataFrame(label="Questions & Answers (Preview)", headers=tests_df_headers) answers_btn.click( several_rag_answer, inputs=[ @@ -488,12 +482,8 @@ def several_rag_answer( ], outputs=[qa_dataframe, gr.File(label="Download Answered File", min_width=40)], ) - questions_file.change( - read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count] - ) - answer_max_line_count.change( - change_showing_excel, answer_max_line_count, qa_dataframe - ) + questions_file.change(read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count]) + answer_max_line_count.change(change_showing_excel, answer_max_line_count, qa_dataframe) return ( inp, answer_prompt_input, diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 04aef1c77..8fbb01c25 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -17,18 +17,23 @@ import json import os -from datetime import datetime from dataclasses import dataclass -from typing import Any, Tuple, Dict, Literal, Optional, List +from datetime import datetime +from typing import Any, Dict, List, Literal, Optional, Tuple import gradio as gr import pandas as pd -from hugegraph_llm.config import prompt, resource_path, huge_settings +from hugegraph_llm.config import huge_settings, index_settings, prompt, resource_path +from hugegraph_llm.models.embeddings.init_embedding import Embeddings +from hugegraph_llm.models.llms.init_llm import LLMs +from hugegraph_llm.operators.graph_rag_task import RAGPipeline +from hugegraph_llm.operators.gremlin_generate_task import GremlinGenerator +from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager from hugegraph_llm.utils.embedding_utils import get_index_folder_name from hugegraph_llm.utils.hugegraph_utils import run_gremlin_query from hugegraph_llm.utils.log import log -from hugegraph_llm.flows.scheduler import SchedulerSingleton +from hugegraph_llm.utils.vector_index_utils import get_vector_index_class @dataclass @@ -81,12 +86,8 @@ def store_schema(schema, question, gremlin_prompt): def build_example_vector_index(temp_file) -> dict: - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - index_path = os.path.join(resource_path, folder_name, "gremlin_examples") - if not os.path.exists(index_path): - os.makedirs(index_path) + vector_index = get_vector_index_class(index_settings.cur_vector_index) + assert vector_index, "vector db name is error" if temp_file is None: full_path = os.path.join(resource_path, "demo", "text2gremlin.csv") else: @@ -95,12 +96,12 @@ def build_example_vector_index(temp_file) -> dict: timestamp = datetime.now().strftime("%Y%m%d%H%M%S") _, file_name = os.path.split(f"{name}_{timestamp}{ext}") log.info("Copying file to: %s", file_name) - target_file = os.path.join( - resource_path, folder_name, "gremlin_examples", file_name - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + target_file = os.path.join(resource_path, folder_name, "gremlin_examples", file_name) try: import shutil + os.makedirs(os.path.dirname(target_file), exist_ok=True) shutil.copy2(full_path, target_file) log.info("Successfully copied file to: %s", target_file) except (OSError, IOError) as e: @@ -115,10 +116,11 @@ def build_example_vector_index(temp_file) -> dict: else: log.critical("Unsupported file format. Please input a JSON or CSV file.") return {"error": "Unsupported file format. Please input a JSON or CSV file."} - - return SchedulerSingleton.get_instance().schedule_flow( - "build_examples_index", examples + builder = GremlinGenerator( + llm=LLMs().get_text2gql_llm(), + embedding=Embeddings().get_embedding(), ) + return builder.example_index_build(examples, vector_index=vector_index).run() def _process_schema(schema, generator, sm): @@ -180,18 +182,56 @@ def _execute_queries(context, output_types): context["raw_exec_res"] = "" +def gremlin_generate( + inp, + example_num, + schema, + gremlin_prompt, + requested_outputs: Optional[List[str]] = None, +) -> GremlinResult: + vector_index = get_vector_index_class(index_settings.cur_vector_index) + generator = GremlinGenerator( + llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() + ) + sm = SchemaManager(graph_name=schema) + + processed_schema, short_schema = _process_schema(schema, generator, sm) + if processed_schema is None and short_schema is None: + return GremlinResult.error("Invalid JSON schema, please check the format carefully.") + + updated_schema = sm.simple_schema(processed_schema) if short_schema else processed_schema + store_schema(str(updated_schema), inp, gremlin_prompt) + + output_types = _configure_output_types(requested_outputs) + + context = ( + generator.example_index_query(example_num, vector_index) + .gremlin_generate_synthesize(updated_schema, gremlin_prompt) + .run(query=inp) + ) + + _execute_queries(context, output_types) + + match_result = json.dumps( + context.get("match_result", "No Results"), ensure_ascii=False, indent=2 + ) + return GremlinResult.success_result( + match_result=match_result, + template_gremlin=context["result"], + raw_gremlin=context["raw_result"], + template_exec=context["template_exec_res"], + raw_exec=context["raw_exec_res"], + ) + + def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: - mini_schema = {} + mini_schema = {} # type: ignore # Add necessary vertexlabels items (3) if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = { - key: vertex[key] - for key in ["id", "name", "properties"] - if key in vertex - } + new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) @@ -210,40 +250,17 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: def gremlin_generate_for_ui(inp, example_num, schema, gremlin_prompt): """UI wrapper for gremlin_generate that returns tuple for Gradio compatibility""" - # Execute via scheduler - try: - res = SchedulerSingleton.get_instance().schedule_flow( - "text2gremlin", - inp, - int(example_num) if isinstance(example_num, (int, float, str)) else 2, - schema, - gremlin_prompt, - [ - "match_result", - "template_gremlin", - "raw_gremlin", - "template_execution_result", - "raw_execution_result", - ], - ) - except Exception as e: # pylint: disable=broad-except - log.error("UI text2gremlin error: %s", e) - return json.dumps({"error": str(e)}, ensure_ascii=False), "", "", "", "" - - # Backward-compatible mapping for outputs - match_result = res.get("match_result", []) - match_result_str = ( - json.dumps(match_result, ensure_ascii=False, indent=2) - if isinstance(match_result, (list, dict)) - else str(match_result) - ) + result = gremlin_generate(inp, example_num, schema, gremlin_prompt) + + if not result.success: + return result.match_result, "", "", "", "" return ( - match_result_str, - res.get("template_gremlin", "") or "", - res.get("raw_gremlin", "") or "", - res.get("template_execution_result", "") or "", - res.get("raw_execution_result", "") or "", + result.match_result, + result.template_gremlin or "", + result.raw_gremlin or "", + result.template_exec_result or "", + result.raw_exec_result or "", ) @@ -264,6 +281,7 @@ def create_text2gremlin_block() -> Tuple: out = gr.Textbox(label="Result Message") with gr.Row(): btn = gr.Button("Build Example Vector Index", variant="primary") + btn.click(build_example_vector_index, inputs=[file], outputs=[out]) # pylint: disable=no-member gr.Markdown("## Nature Language To Gremlin") @@ -279,12 +297,8 @@ def create_text2gremlin_block() -> Tuple: language="javascript", elem_classes="code-container-show", ) - initialized_out = gr.Textbox( - label="Gremlin With Template", show_copy_button=True - ) - raw_out = gr.Textbox( - label="Gremlin Without Template", show_copy_button=True - ) + initialized_out = gr.Textbox(label="Gremlin With Template", show_copy_button=True) + raw_out = gr.Textbox(label="Gremlin Without Template", show_copy_button=True) tmpl_exec_out = gr.Code( label="Query With Template Output", language="json", @@ -336,21 +350,25 @@ def graph_rag_recall( get_vertex_only: bool = False, ) -> dict: store_schema(prompt.text2gql_graph_schema, query, gremlin_prompt) - context = SchedulerSingleton.get_instance().schedule_flow( - "rag_graph_only", - query=query, - gremlin_tmpl_num=gremlin_tmpl_num, - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - custom_related_information=custom_related_information, - gremlin_prompt=gremlin_prompt, - max_graph_items=max_graph_items, - topk_return_results=topk_return_results, + rag = RAGPipeline() + rag.extract_keywords().keywords_to_vid( + vector_index_str=index_settings.cur_vector_index, vector_dis_threshold=vector_dis_threshold, topk_per_keyword=topk_per_keyword, - is_graph_rag_recall=True, - is_vector_only=get_vertex_only, ) + + if not get_vertex_only: + rag.import_schema(huge_settings.graph_name).query_graphdb( + num_gremlin_generate_example=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, + max_graph_items=max_graph_items, + ).merge_dedup_rerank( + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + topk_return_results=topk_return_results, + ) + context = rag.run(verbose=True, query=query, graph_search=True) return context @@ -361,13 +379,45 @@ def gremlin_generate_selective( gremlin_prompt_input: str, requested_outputs: Optional[List[str]] = None, ) -> Dict[str, Any]: - response_dict = SchedulerSingleton.get_instance().schedule_flow( - "text2gremlin", - inp, - example_num, - schema_input, - gremlin_prompt_input, - requested_outputs, + """ + Wraps the gremlin_generate function to return a dictionary of outputs + based on the requested_outputs list of strings. + """ + output_keys = [ + "match_result", + "template_gremlin", + "raw_gremlin", + "template_execution_result", + "raw_execution_result", + ] + if not requested_outputs: # None or empty list + requested_outputs = output_keys + + result = gremlin_generate( + inp, example_num, schema_input, gremlin_prompt_input, requested_outputs ) - return response_dict + outputs_dict: Dict[str, Any] = {} + + if not result.success: + # Handle error case + if "match_result" in requested_outputs: + outputs_dict["match_result"] = result.match_result + if result.error_message: + outputs_dict["error_detail"] = result.error_message + return outputs_dict + + # Handle successful case + output_mapping = { + "match_result": result.match_result, + "template_gremlin": result.template_gremlin, + "raw_gremlin": result.raw_gremlin, + "template_execution_result": result.template_exec_result, + "raw_execution_result": result.raw_exec_result, + } + + for key in requested_outputs: + if key in output_mapping: + outputs_dict[key] = output_mapping[key] + + return outputs_dict diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 56b5de4b3..5f63e4964 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -23,25 +23,23 @@ import gradio as gr -from hugegraph_llm.config import huge_settings -from hugegraph_llm.config import prompt -from hugegraph_llm.config import resource_path -from hugegraph_llm.flows.scheduler import SchedulerSingleton +from hugegraph_llm.config import huge_settings, prompt, resource_path +from hugegraph_llm.models.llms.init_llm import LLMs +from hugegraph_llm.operators.llm_op.prompt_generate import PromptGenerate from hugegraph_llm.utils.graph_index_utils import ( - get_graph_index_info, - clean_all_graph_index, + build_schema, clean_all_graph_data, - update_vid_embedding, + clean_all_graph_index, extract_graph, + get_graph_index_info, import_graph_data, - build_schema, + update_vid_embedding, ) from hugegraph_llm.utils.hugegraph_utils import check_graph_db_connection from hugegraph_llm.utils.log import log from hugegraph_llm.utils.vector_index_utils import ( - clean_vector_index, build_vector_index, - get_vector_index_info, + clean_vector_index, ) @@ -60,17 +58,25 @@ def store_prompt(doc, schema, example_prompt): def generate_prompt_for_ui(source_text, scenario, example_name): """ - Handles the UI logic for generating a new prompt using the new workflow architecture. + Handles the UI logic for generating a new prompt. It calls the PromptGenerate operator. """ if not all([source_text, scenario, example_name]): gr.Warning("Please provide original text, expected scenario, and select an example!") return gr.update() try: - # using new architecture - scheduler = SchedulerSingleton.get_instance() - result = scheduler.schedule_flow("prompt_generate", source_text, scenario, example_name) + prompt_generator = PromptGenerate(llm=LLMs().get_chat_llm()) + context = { + "source_text": source_text, + "scenario": scenario, + "example_name": example_name, + } + result_context = prompt_generator.run(context) + # Presents the result of generating prompt + generated_prompt = result_context.get( + "generated_extract_prompt", "Generation failed. Please check the logs." + ) gr.Info("Prompt generated successfully!") - return result + return generated_prompt except Exception as e: log.error("Error generating Prompt: %s", e, exc_info=True) raise gr.Error(f"Error generating Prompt: {e}") from e @@ -292,8 +298,7 @@ def create_vector_graph_block(): gr.Markdown("---") with gr.Accordion("Graph Schema Generator", open=False): gr.Markdown( - "Provide **query examples** and **few-shot examples**, " - "then click **Generate Schema** to automatically create graph schema." + "Provide **query examples** and **few-shot examples**, then click **Generate Schema** to automatically create graph schema." ) with gr.Row(): query_example = gr.Code( @@ -311,9 +316,10 @@ def create_vector_graph_block(): max_lines=15, ) build_schema_bt = gr.Button("Generate Schema", variant="primary") + _create_prompt_helper_block(demo, input_text, info_extract_template) - vector_index_btn0.click(get_vector_index_info, outputs=out).then( + vector_index_btn0.click(get_graph_index_info, outputs=out).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) @@ -344,7 +350,6 @@ def create_vector_graph_block(): inputs=[input_text, input_schema, info_extract_template], ) - # origin_out = gr.Textbox(visible=False) graph_extract_bt.click( extract_graph, inputs=[input_file, input_text, input_schema, info_extract_template], @@ -361,18 +366,13 @@ def create_vector_graph_block(): inputs=[input_text, input_schema, info_extract_template], ) - # TODO: we should store the examples after the user changed them. build_schema_bt.click( _build_schema_and_provide_feedback, inputs=[input_text, query_example, few_shot], outputs=[input_schema], ).then( store_prompt, - inputs=[ - input_text, - input_schema, - info_extract_template, - ], # TODO: Store the updated examples + inputs=[input_text, input_schema, info_extract_template], ) def on_tab_select(input_f, input_t, evt: gr.SelectData): diff --git a/hugegraph-llm/src/hugegraph_llm/flows/common.py b/hugegraph-llm/src/hugegraph_llm/flows/common.py index e2348466c..19a6ca2d4 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/common.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/common.py @@ -46,9 +46,7 @@ def post_deal(self, *args, **kwargs): """ pass - async def post_deal_stream( - self, pipeline=None - ) -> AsyncGenerator[Dict[str, Any], None]: + async def post_deal_stream(self, pipeline=None) -> AsyncGenerator[Dict[str, Any], None]: """ Streaming post-processing interface. Subclasses can override this method as needed. @@ -59,7 +57,7 @@ async def post_deal_stream( return try: state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - log.info(f"{flow_name} post processing success") + log.info("%s post processing success", flow_name) stream_flow = state_json.get("stream_generator") if stream_flow is None: yield {"error": "No stream_generator found in workflow state"} @@ -67,5 +65,5 @@ async def post_deal_stream( async for chunk in stream_flow: yield chunk except Exception as e: - log.error(f"{flow_name} post processing failed: {e}") + log.error("%s post processing failed: %s", flow_name, e) yield {"error": f"Post processing failed: {str(e)}"} diff --git a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py index 439c3a346..a5043bdc7 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py @@ -14,19 +14,15 @@ # limitations under the License. import json -import os -from hugegraph_llm.config import huge_settings, llm_settings, resource_path +from hugegraph_llm.config import huge_settings from hugegraph_llm.flows.common import BaseFlow -from hugegraph_llm.indices.vector_index import VectorIndex -from hugegraph_llm.models.embeddings.init_embedding import model_map +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState from hugegraph_llm.nodes.hugegraph_node.fetch_graph_data import FetchGraphDataNode from PyCGraph import GPipeline -from hugegraph_llm.utils.embedding_utils import ( - get_filename_prefix, - get_index_folder_name, -) +from hugegraph_llm.utils.embedding_utils import get_index_folder_name class GetGraphIndexInfoFlow(BaseFlow): @@ -48,21 +44,15 @@ def build_flow(self, *args, **kwargs): def post_deal(self, pipeline=None): graph_summary_info = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, - model_map.get(llm_settings.embedding_type, None), - ) - try: - vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) - except (RuntimeError, OSError): + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + if not FaissVectorIndex.exist(folder_name, "graph_vids"): return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) + embed_dim = Embeddings().get_embedding().get_embedding_dim() + vector_index = FaissVectorIndex.from_name(embed_dim, folder_name, "graph_vids") + vector_index_info = vector_index.get_vector_index_info() graph_summary_info["vid_index"] = { - "embed_dim": vector_index.index.d, - "num_vectors": vector_index.index.ntotal, - "num_vids": len(vector_index.properties), + "embed_dim": vector_index_info["embed_dim"], + "num_vectors": vector_index_info["vector_info"]["chunk_vector_num"], + "num_vids": vector_index_info["vector_info"]["graph_properties_vector_num"], } return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py index c4cfd46dc..25b901d47 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py @@ -85,12 +85,8 @@ def prepare( prepared_input.graph_vector_answer = graph_vector_answer prepared_input.gremlin_tmpl_num = gremlin_tmpl_num prepared_input.gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt - prepared_input.max_graph_items = ( - max_graph_items or huge_settings.max_graph_items - ) - prepared_input.topk_per_keyword = ( - topk_per_keyword or huge_settings.topk_per_keyword - ) + prepared_input.max_graph_items = max_graph_items or huge_settings.max_graph_items + prepared_input.topk_per_keyword = topk_per_keyword or huge_settings.topk_per_keyword prepared_input.topk_return_results = ( topk_return_results or huge_settings.topk_return_results ) @@ -175,5 +171,5 @@ def post_deal(self, pipeline=None): else res ) except Exception as e: - log.error(f"RAGGraphOnlyFlow post processing failed: {e}") + log.error("RAGGraphOnlyFlow post processing failed: %s", e) return {"error": f"Post processing failed: {str(e)}"} diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py index 9fd4d96e0..2e4e69826 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py @@ -70,15 +70,11 @@ def prepare( prepared_input.graph_ratio = graph_ratio prepared_input.gremlin_tmpl_num = gremlin_tmpl_num prepared_input.gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt - prepared_input.max_graph_items = ( - max_graph_items or huge_settings.max_graph_items - ) + prepared_input.max_graph_items = max_graph_items or huge_settings.max_graph_items prepared_input.topk_return_results = ( topk_return_results or huge_settings.topk_return_results ) - prepared_input.topk_per_keyword = ( - topk_per_keyword or huge_settings.topk_per_keyword - ) + prepared_input.topk_per_keyword = topk_per_keyword or huge_settings.topk_per_keyword prepared_input.vector_dis_threshold = ( vector_dis_threshold or huge_settings.vector_dis_threshold ) @@ -118,19 +114,11 @@ def build_flow(self, **kwargs): # Register nodes and their dependencies pipeline.registerGElement(vector_query_node, set(), "vector") pipeline.registerGElement(keyword_extract_node, set(), "keyword") - pipeline.registerGElement( - semantic_id_query_node, {keyword_extract_node}, "semantic" - ) + pipeline.registerGElement(semantic_id_query_node, {keyword_extract_node}, "semantic") pipeline.registerGElement(schema_node, set(), "schema") - pipeline.registerGElement( - graph_query_node, {schema_node, semantic_id_query_node}, "graph" - ) - pipeline.registerGElement( - merge_rerank_node, {graph_query_node, vector_query_node}, "merge" - ) - pipeline.registerGElement( - answer_synthesize_node, {merge_rerank_node}, "graph_vector" - ) + pipeline.registerGElement(graph_query_node, {schema_node, semantic_id_query_node}, "graph") + pipeline.registerGElement(merge_rerank_node, {graph_query_node, vector_query_node}, "merge") + pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "graph_vector") log.info("RAGGraphVectorFlow pipeline built successfully") return pipeline @@ -147,5 +135,5 @@ def post_deal(self, pipeline=None): "graph_vector_answer": res.get("graph_vector_answer", ""), } except Exception as e: - log.error(f"RAGGraphVectorFlow post processing failed: {e}") + log.error("RAGGraphVectorFlow post processing failed: %s", e) return {"error": f"Post processing failed: {str(e)}"} diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py index f62e574bb..d328d6c6e 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py @@ -78,9 +78,7 @@ def build_flow(self, **kwargs): def post_deal(self, pipeline=None): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) + return json.dumps({"error": "No pipeline provided"}, ensure_ascii=False, indent=2) try: res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() log.info("RAGRawFlow post processing success") @@ -91,7 +89,7 @@ def post_deal(self, pipeline=None): "graph_vector_answer": res.get("graph_vector_answer", ""), } except Exception as e: - log.error(f"RAGRawFlow post processing failed: {e}") + log.error("RAGRawFlow post processing failed: %s", e) return json.dumps( {"error": f"Post processing failed: {str(e)}"}, ensure_ascii=False, diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py index c727eacce..0766903b1 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py @@ -93,18 +93,14 @@ def build_flow(self, **kwargs): # Register nodes and dependencies, keep naming consistent with original pipeline.registerGElement(only_vector_query_node, set(), "only_vector") - pipeline.registerGElement( - merge_rerank_node, {only_vector_query_node}, "merge_two" - ) + pipeline.registerGElement(merge_rerank_node, {only_vector_query_node}, "merge_two") pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "vector") log.info("RAGVectorOnlyFlow pipeline built successfully") return pipeline def post_deal(self, pipeline=None): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) + return json.dumps({"error": "No pipeline provided"}, ensure_ascii=False, indent=2) try: res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() log.info("RAGVectorOnlyFlow post processing success") @@ -115,7 +111,7 @@ def post_deal(self, pipeline=None): "graph_vector_answer": res.get("graph_vector_answer", ""), } except Exception as e: - log.error(f"RAGVectorOnlyFlow post processing failed: {e}") + log.error("RAGVectorOnlyFlow post processing failed: %s", e) return json.dumps( {"error": f"Post processing failed: {str(e)}"}, ensure_ascii=False, diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index 388c8c655..f5904bb7a 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -165,9 +165,7 @@ async def schedule_stream_flow(self, flow_name: str, *args, **kwargs): else: try: # fetch pipeline & prepare input for flow - prepared_input: WkFlowInput = pipeline.getGParamWithNoEmpty( - "wkflow_input" - ) + prepared_input: WkFlowInput = pipeline.getGParamWithNoEmpty("wkflow_input") prepared_input.stream = True flow.prepare(prepared_input, *args, **kwargs) status = pipeline.run() diff --git a/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py b/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py index 694ca014d..31b3d21f3 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py @@ -17,6 +17,7 @@ from typing import Optional + from pyhugegraph.client import PyHugeClient from hugegraph_llm.config import huge_settings diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py deleted file mode 100644 index f85483185..000000000 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py +++ /dev/null @@ -1,170 +0,0 @@ -# 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 os -import pickle as pkl -from copy import deepcopy -from typing import List, Any, Set, Union - -import faiss -import numpy as np - -from hugegraph_llm.utils.log import log - -INDEX_FILE_NAME = "index.faiss" -PROPERTIES_FILE_NAME = "properties.pkl" - - -class VectorIndex: - """Comment""" - - def __init__(self, embed_dim: int = 1024): - self.index = faiss.IndexFlatL2(embed_dim) - self.properties = [] - - @staticmethod - def from_index_file( - dir_path: str, filename_prefix: str = None, record_miss: bool = True - ) -> "VectorIndex": - """Load index from files, supporting model-specific filenames. - - This method loads a Faiss index and its corresponding properties from a directory. - It handles model-specific filenames by constructing them inline using f-strings. - If the specified files are not found, it returns a new, empty VectorIndex instance. - It also performs a consistency check to ensure the number of vectors in the index - matches the number of properties. - """ - index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = ( - f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME - ) - index_file = os.path.join(dir_path, index_name) - properties_file = os.path.join(dir_path, property_name) - miss_files = [f for f in [index_file, properties_file] if not os.path.exists(f)] - if miss_files: - if record_miss: - log.warning( - "Missing vector files: %s. \nNeed create a new one for it.", - ", ".join(miss_files), - ) - return VectorIndex() - - try: - faiss_index = faiss.read_index(index_file) - with open(properties_file, "rb") as f: - properties = pkl.load(f) - except (RuntimeError, pkl.UnpicklingError, OSError) as e: - log.error( - "Failed to load index files for model '%s': %s", filename_prefix or "default", e - ) - raise RuntimeError( - f"Could not load index files for model '{filename_prefix or 'default'}'. " - f"Original error ({type(e).__name__}): {e}" - ) from e - - if faiss_index.ntotal != len(properties): - raise RuntimeError( - f"Data inconsistency: index for model '{filename_prefix or 'default'}' has " - f"{faiss_index.ntotal} vectors, but {len(properties)} properties." - ) - - embed_dim = faiss_index.d - vector_index = VectorIndex(embed_dim) - vector_index.index = faiss_index - vector_index.properties = properties - return vector_index - - def to_index_file(self, dir_path: str, filename_prefix: str = None): - """Save index to files, supporting model-specific filenames.""" - if not os.path.exists(dir_path): - os.makedirs(dir_path) - - index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = ( - f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME - ) - index_file = os.path.join(dir_path, index_name) - properties_file = os.path.join(dir_path, property_name) - faiss.write_index(self.index, index_file) - with open(properties_file, "wb") as f: - pkl.dump(self.properties, f) - - def add(self, vectors: List[List[float]], props: List[Any]): - if len(vectors) == 0: - return - - if self.index.ntotal == 0 and len(vectors[0]) != self.index.d: - self.index = faiss.IndexFlatL2(len(vectors[0])) - self.index.add(np.array(vectors)) - self.properties.extend(props) - - def remove(self, props: Union[Set[Any], List[Any]]) -> int: - if isinstance(props, list): - props = set(props) - indices = [] - remove_num = 0 - - for i, p in enumerate(self.properties): - if p in props: - indices.append(i) - remove_num += 1 - self.index.remove_ids(np.array(indices)) - self.properties = [p for i, p in enumerate(self.properties) if i not in indices] - return remove_num - - def search( - self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9 - ) -> List[Any]: - if self.index.ntotal == 0: - return [] - - if len(query_vector) != self.index.d: - raise ValueError("Query vector dimension does not match index dimension!") - - distances, indices = self.index.search(np.array([query_vector]), top_k) - results = [] - for dist, i in zip(distances[0], indices[0]): - if dist < dis_threshold: # Smaller distances indicate higher similarity - results.append(deepcopy(self.properties[i])) - log.debug("[✓] Add valid distance %s to results.", dist) - else: - log.debug( - "[x] Distance %s >= threshold %s, ignore this result.", dist, dis_threshold - ) - return results - - @staticmethod - def clean(dir_path: str, filename_prefix: str = None): - """Clean index files, supporting model-specific filenames. - - This method deletes the index and properties files associated with a specific model. - If model_name is None, it targets the default files. - """ - index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = ( - f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME - ) - index_file = os.path.join(dir_path, index_name) - properties_file = os.path.join(dir_path, property_name) - - for file in [index_file, properties_file]: - if os.path.exists(file): - try: - os.remove(file) - log.info("Removed index file: %s", file) - except OSError as e: - log.error("Error removing file %s: %s", file, e) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py new file mode 100644 index 000000000..2e1cfc267 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/base.py @@ -0,0 +1,107 @@ +# 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 abc import ABC, abstractmethod +from typing import Any, Dict, List, Set, Union + + +class VectorStoreBase(ABC): + """ + Abstract base class defining the interface for a vector store. + Implementations must support adding, removing, searching vectors, + saving/loading from disk, and cleaning up resources. + """ + + @abstractmethod + def add(self, vectors: List[List[float]], props: List[Any]): + """ + Add a list of vectors and their corresponding properties to the store. + + Args: + vectors (List[List[float]]): List of embedding vectors. + props (List[Any]): List of associated metadata or properties for each vector. + """ + + @abstractmethod + def get_all_properties(self) -> list[str]: + """ + #TODO: finish comment + """ + + @abstractmethod + def remove(self, props: Union[Set[Any], List[Any]]) -> int: + """ + Remove vectors based on their associated properties. + + Args: + props (Union[Set[Any], List[Any]]): Properties of vectors to remove. + + Returns: + int: Number of vectors removed. + """ + + @abstractmethod + def search( + self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9 + ) -> List[Any]: + """ + Search for the top_k most similar vectors to the query vector. + + Args: + query_vector (List[float]): The vector to query against the index. + top_k (int): Number of top results to return. + dis_threshold (float): Distance threshold below which results are considered relevant. + + Returns: + List[Any]: List of properties of the matched vectors. + """ + + @abstractmethod + def save_index_by_name(self, *name: str): + """ + #TODO: finish comment + """ + + @abstractmethod + def get_vector_index_info( + self, + ) -> Dict: + """ + #TODO: finish comment + """ + + @staticmethod + @abstractmethod + def from_name(embed_dim: int, *name: str) -> "VectorStoreBase": + """ + #TODO: finish comment + """ + + @staticmethod + @abstractmethod + def exist(*name: str) -> bool: + """ + #TODO: finish comment + """ + + @staticmethod + @abstractmethod + def clean(*name: str) -> bool: + """ + #TODO: finish comment + """ diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.py new file mode 100644 index 000000000..a8f23155a --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/faiss_vector_store.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. + +import os +import pickle as pkl +from copy import deepcopy +from typing import Any, Dict, List, Set, Union + +import faiss +import numpy as np + +from hugegraph_llm.config import resource_path +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +from hugegraph_llm.utils.log import log + +INDEX_FILE_NAME = "index.faiss" +PROPERTIES_FILE_NAME = "properties.pkl" + + +class FaissVectorIndex(VectorStoreBase): + def __init__(self, embed_dim: int = 1024): + self.index = faiss.IndexFlatL2(embed_dim) + self.properties: list[Any] = [] + + def save_index_by_name(self, *name: str): + os.makedirs(os.path.join(resource_path, *name), exist_ok=True) + index_file = os.path.join(resource_path, *name, INDEX_FILE_NAME) + properties_file = os.path.join(resource_path, *name, PROPERTIES_FILE_NAME) + faiss.write_index(self.index, index_file) + with open(properties_file, "wb") as f: + pkl.dump(self.properties, f) + + def add(self, vectors: List[List[float]], props: List[Any]): + if len(vectors) == 0: + return + if self.index.ntotal == 0 and len(vectors[0]) != self.index.d: + self.index = faiss.IndexFlatL2(len(vectors[0])) + self.index.add(np.array(vectors)) + self.properties.extend(props) + + def remove(self, props: Union[Set[Any], List[Any]]) -> int: + if isinstance(props, list): + props = set(props) + indices = [] + remove_num = 0 + + for i, p in enumerate(self.properties): + if p in props: + indices.append(i) + remove_num += 1 + self.index.remove_ids(np.array(indices)) + self.properties = [p for i, p in enumerate(self.properties) if i not in indices] + return remove_num + + def search( + self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9 + ) -> List[Any]: + if self.index.ntotal == 0: + return [] + + if len(query_vector) != self.index.d: + raise ValueError("Query vector dimension does not match index dimension!") + + distances, indices = self.index.search(np.array([query_vector]), top_k) + results = [] + for dist, i in zip(distances[0], indices[0]): + if dist < dis_threshold: + results.append(deepcopy(self.properties[i])) + log.debug("[✓] Add valid distance %s to results.", dist) + else: + log.debug( + "[x] Distance %s >= threshold %s, ignore this result.", + dist, + dis_threshold, + ) + return results + + def get_all_properties(self) -> list[Any]: + return self.properties + + def get_vector_index_info( + self, + ) -> Dict: + return { + "embed_dim": self.index.d, + "vector_info": { + "chunk_vector_num": self.index.ntotal, + "graph_vid_vector_num": self.index.ntotal, + "graph_properties_vector_num": len(self.properties), + }, + } + + @staticmethod + def clean(*name: str): + index_file = os.path.join(resource_path, *name, INDEX_FILE_NAME) + properties_file = os.path.join(resource_path, *name, PROPERTIES_FILE_NAME) + if os.path.exists(index_file): + os.remove(index_file) + if os.path.exists(properties_file): + os.remove(properties_file) + + @staticmethod + def from_name(embed_dim: int, *name: str) -> "FaissVectorIndex": + index_file = os.path.join(resource_path, *name, INDEX_FILE_NAME) + properties_file = os.path.join(resource_path, *name, PROPERTIES_FILE_NAME) + if not os.path.exists(index_file) or not os.path.exists(properties_file): + log.warning("No index file found, create a new one.") + return FaissVectorIndex(embed_dim) + + faiss_index = faiss.read_index(index_file) + with open(properties_file, "rb") as f: + properties = pkl.load(f) + vector_index = FaissVectorIndex(embed_dim) + if faiss_index.d == vector_index.index.d: + # when dim same, use old + vector_index.index = faiss_index + vector_index.properties = properties + else: + log.warning("dim is different, create a new one.") + return vector_index + + @staticmethod + def exist(*name: str) -> bool: + index_file = os.path.join(resource_path, *name, INDEX_FILE_NAME) + properties_file = os.path.join(resource_path, *name, PROPERTIES_FILE_NAME) + return os.path.exists(index_file) and os.path.exists(properties_file) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py new file mode 100644 index 000000000..b168b8c39 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -0,0 +1,263 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import json +from typing import Any, List, Set, Union + +from pymilvus import ( + Collection, + CollectionSchema, + DataType, + FieldSchema, + connections, + utility, +) + +from hugegraph_llm.config import index_settings +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +from hugegraph_llm.utils.log import log + +COLLECTION_NAME_PREFIX = "hugegraph_llm_" + + +class MilvusVectorIndex(VectorStoreBase): + def __init__( + self, + name: str, + host: str, + port: int, + user="", + password="", + embed_dim: int = 1024, + ): + self.embed_dim = embed_dim + self.host = host + self.port = port + self.name = COLLECTION_NAME_PREFIX + name + connections.connect(host=host, port=port, user=user, password=password) + + if not utility.has_collection(self.name): + self._create_collection() + else: + # dim is different, recreate + existing_collection = Collection(self.name) + existing_schema = existing_collection.schema + for field in existing_schema.fields: + if field.name == "embedding" and field.params.get("dim"): + existing_dim = int(field.params["dim"]) + if existing_dim != self.embed_dim: + log.debug( + "Milvus collection '%s' dimension mismatch: %d != %d. Recreating.", + self.name, + existing_dim, + self.embed_dim, + ) + utility.drop_collection(self.name) + break + + self.collection = Collection(self.name) + + def _create_collection(self): + """Create a new collection in Milvus.""" + id_field = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True) + vector_field = FieldSchema( + name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embed_dim + ) + property_field = FieldSchema(name="property", dtype=DataType.VARCHAR, max_length=65535) + original_id_field = FieldSchema(name="original_id", dtype=DataType.INT64) + + schema = CollectionSchema( + fields=[id_field, vector_field, property_field, original_id_field], + description="Vector index collection", + ) + + collection = Collection(name=self.name, schema=schema) + + index_params = { + "metric_type": "L2", + "index_type": "IVF_FLAT", + "params": {"nlist": 128}, + } + collection.create_index(field_name="embedding", index_params=index_params) + + def save_index_by_name(self, *name: str): + self.collection.flush() + + def _deserialize_property(self, prop) -> str: + """If input is a string, return as-is. If dict or list, convert to JSON string.""" + if isinstance(prop, str): + return prop + return json.dumps(prop) + + def _serialize_property(self, prop: str): + """If input is a JSON string, parse it. Otherwise, return as-is.""" + try: + return json.loads(prop) + except (json.JSONDecodeError, TypeError): + # a simple string + return prop + + def add(self, vectors: List[List[float]], props: List[Any]): + if len(vectors) == 0: + return + + # Get the current count to use as starting index + count = self.collection.num_entities + entities = [] + + for i, (vector, prop) in enumerate(zip(vectors, props)): + idx = count + i + entities.append( + { + "embedding": vector, + "property": self._deserialize_property(prop), + "original_id": idx, + } + ) + + self.collection.insert(entities) + self.collection.flush() + + def remove(self, props: Union[Set[Any], List[Any]]) -> int: + if isinstance(props, list): + props = set(props) + try: + self.collection.load() + remove_num = 0 + for prop in props: + expr = f'property == "{self._deserialize_property(prop)}"' + res = self.collection.delete(expr) + if hasattr(res, "delete_count"): + remove_num += res.delete_count + if remove_num > 0: + self.collection.flush() + return remove_num + finally: + self.collection.release() + + def search( + self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9 + ) -> List[Any]: + try: + if self.collection.num_entities == 0: + return [] + + self.collection.load() + search_params = {"metric_type": "L2", "params": {"nprobe": 10}} + results = self.collection.search( + data=[query_vector], + anns_field="embedding", + param=search_params, + limit=top_k, + output_fields=["property"], + ) + + ret = [] + for hits in results: + for hit in hits: + if hit.distance < dis_threshold: + prop_str = hit.entity.get("property") + prop = self._serialize_property(prop_str) + ret.append(prop) + log.debug("[✓] Add valid distance %s to results.", hit.distance) + else: + log.debug( + "[x] Distance %s >= threshold %s, ignore this result.", + hit.distance, + dis_threshold, + ) + + return ret + + finally: + self.collection.release() + + def get_all_properties(self) -> list[str]: + if self.collection.num_entities == 0: + return [] + + self.collection.load() + try: + results = self.collection.query( + expr='property != ""', + output_fields=["property"], + ) + + return [self._deserialize_property(item["property"]) for item in results] + + finally: + self.collection.release() + + def get_vector_index_info(self) -> dict: + self.collection.load() + try: + embed_dim = None + for field in self.collection.schema.fields: + if field.name == "embedding" and field.dtype == DataType.FLOAT_VECTOR: + embed_dim = int(field.params["dim"]) + break + + if embed_dim is None: + raise ValueError("Could not determine embedding dimension from schema.") + + properties = self.get_all_properties() + return { + "embed_dim": embed_dim, + "vector_info": { + "chunk_vector_num": self.collection.num_entities, + "graph_vid_vector_num": self.collection.num_entities, + "graph_properties_vector_num": len(properties), + }, + } + finally: + self.collection.release() + + @staticmethod + def clean(*name: str): + name_str = "_".join(name) + connections.connect( + host=index_settings.milvus_host, + port=index_settings.milvus_port, + user=index_settings.milvus_user, + password=index_settings.milvus_password, + ) + if utility.has_collection(COLLECTION_NAME_PREFIX + name_str): + utility.drop_collection(COLLECTION_NAME_PREFIX + name_str) + + @staticmethod + def from_name(embed_dim: int, *name: str) -> "MilvusVectorIndex": + name_str = "_".join(name) + assert index_settings.milvus_host, "Qdrant host is not configured" + return MilvusVectorIndex( + name_str, + host=index_settings.milvus_host, + port=index_settings.milvus_port, + user=index_settings.milvus_user, + password=index_settings.milvus_password, + embed_dim=embed_dim, + ) + + @staticmethod + def exist(*name: str) -> bool: + name_str = "_".join(name) + connections.connect( + host=index_settings.milvus_host, + port=index_settings.milvus_port, + user=index_settings.milvus_user, + password=index_settings.milvus_password, + ) + return utility.has_collection(COLLECTION_NAME_PREFIX + name_str) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py new file mode 100644 index 000000000..4342c90ee --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -0,0 +1,221 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from typing import Any, Dict, List, Set, Union + +from qdrant_client import QdrantClient +from qdrant_client.http import models + +from hugegraph_llm.config import index_settings +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +from hugegraph_llm.utils.log import log + +COLLECTION_NAME_PREFIX = "hugegraph_llm_" + + +class QdrantVectorIndex(VectorStoreBase): + def __init__(self, name: str, host: str, port: int, api_key=None, embed_dim: int = 1024): + self.embed_dim = embed_dim + self.host = host + self.port = port + self.name = COLLECTION_NAME_PREFIX + name + self.client = QdrantClient(host=host, port=port, api_key=api_key) + collections = self.client.get_collections().collections + collection_names = [collection.name for collection in collections] + if self.name not in collection_names: + self._create_collection() + else: + collection_info = self.client.get_collection(self.name) + existing_dim = collection_info.config.params.vectors.size # type: ignore + if existing_dim != self.embed_dim: + log.debug( + "Qdrant collection '%s' dimension mismatch: %d != %d. Recreating.", + self.name, + existing_dim, + self.embed_dim, + ) + self.client.delete_collection(self.name) + self._create_collection() + + def _create_collection(self): + """Create a new collection in Qdrant.""" + self.client.create_collection( + collection_name=self.name, + vectors_config=models.VectorParams( + size=self.embed_dim, distance=models.Distance.COSINE + ), + ) + log.info("Created Qdrant collection '%s'", self.name) + + def save_index_by_name(self, *name: str): + # nothing to do when qdrant + pass + + def add(self, vectors: List[List[float]], props: List[Any]): + if len(vectors) == 0: + return + + points = [] + + for i, (vector, prop) in enumerate(zip(vectors, props)): + points.append( + models.PointStruct( + id=i, + vector=vector, + payload={"property": prop}, + ) + ) + + self.client.upsert(collection_name=self.name, points=points, wait=True) + + def remove(self, props: Union[Set[Any], List[Any]]) -> int: + if isinstance(props, list): + props = set(props) + + remove_num = 0 + + for prop in props: + serialized_prop = prop + search_result = self.client.scroll( + collection_name=self.name, + scroll_filter=models.Filter( + must=[ + models.FieldCondition( + key="property", + match=models.MatchValue(value=serialized_prop), + ) + ] + ), + limit=1000, + ) + if search_result and search_result[0]: + point_ids = [point.id for point in search_result[0]] + + if point_ids: + _ = self.client.delete( + collection_name=self.name, + points_selector=models.PointIdsList(points=point_ids), + wait=True, + ) + remove_num += len(point_ids) + + return remove_num + + def search(self, query_vector: List[float], top_k: int = 5, dis_threshold: float = 0.9): + search_result = self.client.search( + collection_name=self.name, query_vector=query_vector, limit=top_k + ) + + result_properties = [] + + for hit in search_result: + distance = 1.0 - hit.score + if distance < dis_threshold: + if hit.payload is not None: + result_properties.append(hit.payload.get("property")) + log.debug("[✓] Add valid distance %s to results.", distance) + else: + log.debug("[x] Hit payload is None, skipping.") + else: + log.debug( + "[x] Distance %s >= threshold %s, ignore this result.", + distance, + dis_threshold, + ) + + return result_properties + + def get_all_properties(self) -> list[str]: + all_properties = [] + offset = None + page_size = 100 + while True: + scroll_result = self.client.scroll( + collection_name=self.name, + offset=offset, + limit=page_size, + with_payload=True, + with_vectors=False, + ) + + points, next_offset = scroll_result + + for point in points: + payload = point.payload + if payload and "property" in payload: + all_properties.append(payload["property"]) + + if next_offset is None or not points: + break + + offset = next_offset + + return all_properties + + def get_vector_index_info(self) -> Dict: + collection_info = self.client.get_collection(self.name) + points_count = collection_info.points_count + embed_dim = collection_info.config.params.vectors.size # type: ignore + + all_properties = self.get_all_properties() + return { + "embed_dim": embed_dim, + "vector_info": { + "chunk_vector_num": points_count, + "graph_vid_vector_num": points_count, + "graph_properties_vector_num": len(all_properties), + }, + } + + @staticmethod + def clean(*name: str): + name_str = "_".join(name) + client = QdrantClient( + host=index_settings.qdrant_host, + port=index_settings.qdrant_port, + api_key=index_settings.qdrant_api_key, + ) + collections = client.get_collections().collections + collection_names = [collection.name for collection in collections] + name_str = COLLECTION_NAME_PREFIX + name_str + if name_str in collection_names: + client.delete_collection(collection_name=name_str) + + @staticmethod + def from_name(embed_dim: int, *name: str) -> "QdrantVectorIndex": + assert index_settings.qdrant_host, "Qdrant host is not configured" + name_str = "_".join(name) + return QdrantVectorIndex( + name=name_str, + host=index_settings.qdrant_host, + port=index_settings.qdrant_port, + embed_dim=embed_dim, + api_key=index_settings.qdrant_api_key, + ) + + @staticmethod + def exist(*name: str) -> bool: + name_str = "_".join(name) + client = QdrantClient( + host=index_settings.qdrant_host, + port=index_settings.qdrant_port, + api_key=index_settings.qdrant_api_key, + ) + collections = client.get_collections().collections + collection_names = [collection.name for collection in collections] + name_str = COLLECTION_NAME_PREFIX + name_str + return name_str in collection_names diff --git a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py index c73242012..f13e11e6f 100644 --- a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py +++ b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py @@ -46,7 +46,7 @@ async def dispatch(self, request: Request, call_next): "%s - Args: %s, IP: %s, URL: %s", request.method, request.query_params, - request.client.host, + request.client.host, # type: ignore request.url, ) return response diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py index 698b92837..62874319c 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py @@ -60,6 +60,12 @@ class BaseEmbedding(ABC): def get_text_embedding(self, text: str) -> List[float]: """Comment""" + @abstractmethod + def get_embedding_dim( + self, + ) -> int: + """Comment""" + @abstractmethod def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch. @@ -81,24 +87,8 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """ @abstractmethod - async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts in a single batch asynchronously. - - This method should efficiently process multiple texts at once by leveraging - the embedding model's batching capabilities, which is typically more efficient - than processing texts individually. - - Parameters - ---------- - texts : List[str] - A list of text strings to be embedded. - - Returns - ------- - List[List[float]] - A list of embedding vectors, where each vector is a list of floats. - The order of embeddings should match the order of input texts. - """ + async def async_get_text_embedding(self, text: str) -> List[float]: + """Comment""" @staticmethod def similarity( diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index 3ad50b3ec..dfd23ccfe 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -17,40 +17,10 @@ from hugegraph_llm.config import llm_settings -from hugegraph_llm.config import LLMConfig from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding from hugegraph_llm.models.embeddings.openai import OpenAIEmbedding -model_map = { - "openai": llm_settings.openai_embedding_model, - "ollama/local": llm_settings.ollama_embedding_model, - "litellm": llm_settings.litellm_embedding_model, -} - - -def get_embedding(llm_settings: LLMConfig): - if llm_settings.embedding_type == "openai": - return OpenAIEmbedding( - model_name=llm_settings.openai_embedding_model, - api_key=llm_settings.openai_embedding_api_key, - api_base=llm_settings.openai_embedding_api_base, - ) - if llm_settings.embedding_type == "ollama/local": - return OllamaEmbedding( - model_name=llm_settings.ollama_embedding_model, - host=llm_settings.ollama_embedding_host, - port=llm_settings.ollama_embedding_port, - ) - if llm_settings.embedding_type == "litellm": - return LiteLLMEmbedding( - model_name=llm_settings.litellm_embedding_model, - api_key=llm_settings.litellm_embedding_api_key, - api_base=llm_settings.litellm_embedding_api_base, - ) - - raise Exception("embedding type is not supported !") - class Embeddings: def __init__(self): @@ -58,22 +28,27 @@ def __init__(self): def get_embedding(self): if self.embedding_type == "openai": + assert llm_settings.openai_embedding_model_dim, "openai_embedding_model_dim is need" return OpenAIEmbedding( + embedding_dimension=llm_settings.openai_embedding_model_dim, model_name=llm_settings.openai_embedding_model, api_key=llm_settings.openai_embedding_api_key, api_base=llm_settings.openai_embedding_api_base, ) if self.embedding_type == "ollama/local": + assert llm_settings.ollama_embedding_model_dim, "ollama_embedding_model_dim is need" return OllamaEmbedding( - model_name=llm_settings.ollama_embedding_model, + embedding_dimension=llm_settings.ollama_embedding_model_dim, + model=llm_settings.ollama_embedding_model, host=llm_settings.ollama_embedding_host, port=llm_settings.ollama_embedding_port, ) if self.embedding_type == "litellm": return LiteLLMEmbedding( + embedding_dimension=llm_settings.litellm_embedding_model_dim, model_name=llm_settings.litellm_embedding_model, api_key=llm_settings.litellm_embedding_api_key, api_base=llm_settings.litellm_embedding_api_base, - ) + ) # type: ignore raise Exception("embedding type is not supported !") diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py index 9d15daa0a..c5effaacf 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py @@ -17,16 +17,11 @@ from typing import List, Optional +from litellm import APIConnectionError, APIError, RateLimitError, aembedding, embedding +from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential + from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.utils.log import log -from tenacity import ( - retry, - stop_after_attempt, - wait_exponential, - retry_if_exception_type, -) - -from litellm import embedding, RateLimitError, APIError, APIConnectionError, aembedding class LiteLLMEmbedding(BaseEmbedding): @@ -34,13 +29,20 @@ class LiteLLMEmbedding(BaseEmbedding): def __init__( self, + embedding_dimension, api_key: Optional[str] = None, api_base: Optional[str] = None, model_name: str = "openai/text-embedding-3-small", # Can be any embedding model supported by LiteLLM ) -> None: self.api_key = api_key self.api_base = api_base - self.model_name = model_name + self.model = model_name + self.embedding_dimension = embedding_dimension + + def get_embedding_dim( + self, + ) -> int: + return self.embedding_dimension @retry( stop=stop_after_attempt(3), @@ -51,7 +53,7 @@ def get_text_embedding(self, text: str) -> List[float]: """Get embedding for a single text.""" try: response = embedding( - model=self.model_name, + model=self.model, input=text, api_key=self.api_key, api_base=self.api_base, @@ -66,7 +68,7 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts.""" try: response = embedding( - model=self.model_name, + model=self.model, input=texts, api_key=self.api_key, api_base=self.api_base, @@ -77,11 +79,26 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: log.error("Error in LiteLLM batch embedding call: %s", e) raise - async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: + async def async_get_text_embedding(self, text: str) -> List[float]: """Get embedding for a single text asynchronously.""" try: response = await aembedding( - model=self.model_name, + model=self.model, + input=text, + api_key=self.api_key, + api_base=self.api_base, + ) + log.info("Token usage: %s", response.usage) + return response.data[0]["embedding"] + except (RateLimitError, APIConnectionError, APIError) as e: + log.error("Error in async LiteLLM embedding call: %s", e) + raise + + async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: + """Get embeddings for multiple texts asynchronously.""" + try: + response = await aembedding( + model=self.model, input=texts, api_key=self.api_key, api_base=self.api_base, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py index a4a8bb098..c02590695 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py @@ -23,19 +23,40 @@ class OllamaEmbedding(BaseEmbedding): - def __init__(self, model_name: str, host: str = "127.0.0.1", port: int = 11434, **kwargs): - self.model_name = model_name + def __init__( + self, + model: str = "quentinz/bge-large-zh-v1.5", + embedding_dimension: int = 1024, + host: str = "127.0.0.1", + port: int = 11434, + **kwargs, + ): + self.model = model self.client = ollama.Client(host=f"http://{host}:{port}", **kwargs) self.async_client = ollama.AsyncClient(host=f"http://{host}:{port}", **kwargs) - self.embedding_dimension = None + self.embedding_dimension = embedding_dimension + + def get_embedding_dim( + self, + ) -> int: + return self.embedding_dimension def get_text_embedding(self, text: str) -> List[float]: - """Get embedding for a single text.""" - return self.get_texts_embeddings([text])[0] + """Comment""" + return list(self.client.embed(model=self.model, input=text)["embeddings"][0]) def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch. + This method efficiently processes multiple texts at once by leveraging + Ollama's batching capabilities, which is more efficient than processing + texts individually. + + Parameters + ---------- + texts : List[str] + A list of text strings to be embedded. + Returns ------- List[List[float]] @@ -49,23 +70,18 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: ) raise AttributeError(error_message) - response = self.client.embed(model=self.model_name, input=texts)["embeddings"] + response = self.client.embed(model=self.model, input=texts)["embeddings"] return [list(inner_sequence) for inner_sequence in response] - async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts in a single batch asynchronously. + async def async_get_text_embedding(self, text: str) -> List[float]: + """Get embedding for a single text asynchronously.""" + response = await self.async_client.embeddings(model=self.model, prompt=text) + return list(response["embedding"]) - Returns - ------- - List[List[float]] - A list of embedding vectors, where each vector is a list of floats. - The order of embeddings matches the order of input texts. - """ - if not hasattr(self.client, "embed"): - error_message = ( - "The required 'embed' method was not found on the Ollama client. " - "Please ensure your ollama library is up-to-date and supports batch embedding. " - ) - raise AttributeError(error_message) - response = await self.async_client.embed(model=self.model_name, input=texts) - return [list(inner_sequence) for inner_sequence in response["embeddings"]] + async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: + # Ollama python client may not provide batch async embeddings; fallback per item + results: List[List[float]] = [] + for t in texts: + response = await self.async_client.embeddings(model=self.model, prompt=t) + results.append(list(response["embedding"])) + return results diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index d0e15f000..135f71000 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -19,11 +19,13 @@ from typing import Optional, List from openai import OpenAI, AsyncOpenAI +from hugegraph_llm.models.embeddings.base import BaseEmbedding -class OpenAIEmbedding: +class OpenAIEmbedding(BaseEmbedding): def __init__( self, + embedding_dimension: int = 1536, model_name: str = "text-embedding-3-small", api_key: Optional[str] = None, api_base: Optional[str] = None, @@ -31,11 +33,17 @@ def __init__( api_key = api_key or "" self.client = OpenAI(api_key=api_key, base_url=api_base) self.aclient = AsyncOpenAI(api_key=api_key, base_url=api_base) - self.model_name = model_name + self.model = model_name + self.embedding_dimension = embedding_dimension + + def get_embedding_dim( + self, + ) -> int: + return self.embedding_dimension def get_text_embedding(self, text: str) -> List[float]: """Comment""" - response = self.client.embeddings.create(input=text, model=self.model_name) + response = self.client.embeddings.create(input=text, model=self.model) return response.data[0].embedding def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: @@ -56,7 +64,7 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: A list of embedding vectors, where each vector is a list of floats. The order of embeddings matches the order of input texts. """ - response = self.client.embeddings.create(input=texts, model=self.model_name) + response = self.client.embeddings.create(input=texts, model=self.model) return [data.embedding for data in response.data] async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: @@ -77,5 +85,9 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] A list of embedding vectors, where each vector is a list of floats. The order of embeddings should match the order of input texts. """ - response = await self.aclient.embeddings.create(input=texts, model=self.model_name) + response = await self.aclient.embeddings.create(input=texts, model=self.model) return [data.embedding for data in response.data] + + async def async_get_text_embedding(self, text: str) -> List[float]: + response = await self.aclient.embeddings.create(input=[text], model=self.model) + return response.data[0].embedding diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py index 78f53e231..d29ec2c4e 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py @@ -16,8 +16,8 @@ from typing import Dict, Any from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.common_op.merge_dedup_rerank import MergeDedupRerank -from hugegraph_llm.models.embeddings.init_embedding import get_embedding -from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.models.embeddings.init_embedding import Embeddings +from hugegraph_llm.config import huge_settings from hugegraph_llm.utils.log import log @@ -34,7 +34,7 @@ def node_init(self): """ try: # Read user configuration parameters from wk_input - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() graph_ratio = self.wk_input.graph_ratio or 0.5 rerank_method = self.wk_input.rerank_method or "bleu" near_neighbor_first = self.wk_input.near_neighbor_first or False @@ -53,7 +53,7 @@ def node_init(self): ) return super().node_init() except Exception as e: - log.error(f"Failed to initialize MergeRerankNode: {e}") + log.error("Failed to initialize MergeRerankNode: %s", e) from PyCGraph import CStatus return CStatus(-1, f"MergeRerankNode initialization failed: {e}") @@ -72,12 +72,14 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: merged_count = len(result.get("merged_result", [])) log.info( - f"Merge and rerank completed: {vector_count} vector results, " - f"{graph_count} graph results, {merged_count} merged results" + "Merge and rerank completed: %d vector results, %d graph results, %d merged results", + vector_count, + graph_count, + merged_count, ) return result except Exception as e: - log.error(f"Merge and rerank failed: {e}") + log.error("Merge and rerank failed: %s", e) return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py index 7bc9dab69..233e73e8c 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py @@ -13,63 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json +from typing import Dict, Any from PyCGraph import CStatus -from typing import Dict, Any, Tuple, List, Set, Optional from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery from hugegraph_llm.config import huge_settings, prompt -from hugegraph_llm.operators.operator_list import OperatorList from hugegraph_llm.utils.log import log -from pyhugegraph.client import PyHugeClient - -# TODO: remove 'as('subj)' step -VERTEX_QUERY_TPL = "g.V({keywords}).limit(8).as('subj').toList()" - -# TODO: we could use a simpler query (like kneighbor-api to get the edges) -# TODO: test with profile()/explain() to speed up the query -VID_QUERY_NEIGHBOR_TPL = """\ -g.V({keywords}) -.repeat( - bothE({edge_labels}).limit({edge_limit}).otherV().dedup() -).times({max_deep}).emit() -.simplePath() -.path() -.by(project('label', 'id', 'props') - .by(label()) - .by(id()) - .by(valueMap().by(unfold())) -) -.by(project('label', 'inV', 'outV', 'props') - .by(label()) - .by(inV().id()) - .by(outV().id()) - .by(valueMap().by(unfold())) -) -.limit({max_items}) -.toList() -""" - -PROPERTY_QUERY_NEIGHBOR_TPL = """\ -g.V().has('{prop}', within({keywords})) -.repeat( - bothE({edge_labels}).limit({edge_limit}).otherV().dedup() -).times({max_deep}).emit() -.simplePath() -.path() -.by(project('label', 'props') - .by(label()) - .by(valueMap().by(unfold())) -) -.by(project('label', 'inV', 'outV', 'props') - .by(label()) - .by(inV().values('{prop}')) - .by(outV().values('{prop}')) - .by(valueMap().by(unfold())) -) -.limit({max_items}) -.toList() -""" class GraphQueryNode(BaseNode): @@ -77,387 +27,42 @@ class GraphQueryNode(BaseNode): Graph query node, responsible for retrieving relevant information from the graph database. """ + graph_rag_query: GraphRAGQuery + def node_init(self): """ Initialize the graph query operator. """ try: - self._client: PyHugeClient = 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, - ) - self._max_deep = self.wk_input.max_deep or 2 - self._max_items = ( - self.wk_input.max_graph_items or huge_settings.max_graph_items - ) - self._prop_to_match = self.wk_input.prop_to_match - self._num_gremlin_generate_example = self.wk_input.gremlin_tmpl_num or -1 - self.gremlin_prompt = ( - self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt + graph_name = huge_settings.graph_name + if not graph_name: + return CStatus(-1, "graph_name is required in wk_input") + + max_deep = self.wk_input.max_deep or 2 + max_graph_items = self.wk_input.max_graph_items or huge_settings.max_graph_items + max_v_prop_len = self.wk_input.max_v_prop_len or 2048 + max_e_prop_len = self.wk_input.max_e_prop_len or 256 + prop_to_match = self.wk_input.prop_to_match + num_gremlin_generate_example = self.wk_input.gremlin_tmpl_num or -1 + gremlin_prompt = self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt + + # Initialize GraphRAGQuery operator + self.graph_rag_query = GraphRAGQuery( + max_deep=max_deep, + max_graph_items=max_graph_items, + max_v_prop_len=max_v_prop_len, + max_e_prop_len=max_e_prop_len, + prop_to_match=prop_to_match, + num_gremlin_generate_example=num_gremlin_generate_example, + gremlin_prompt=gremlin_prompt, ) - self._limit_property = huge_settings.limit_property.lower() == "true" - self._max_v_prop_len = self.wk_input.max_v_prop_len or 2048 - self._max_e_prop_len = self.wk_input.max_e_prop_len or 256 - self._schema = "" - self.operator_list = OperatorList(None, None) return super().node_init() except Exception as e: - log.error(f"Failed to initialize GraphQueryNode: {e}") + log.error("Failed to initialize GraphQueryNode: %s", e) return CStatus(-1, f"GraphQueryNode initialization failed: {e}") - # TODO: move this method to a util file for reuse (remove self param) - def init_client(self, context): - """Initialize the HugeGraph client from context or default settings.""" - # pylint: disable=R0915 (too-many-statements) - if self._client is None: - if isinstance(context.get("graph_client"), PyHugeClient): - self._client = context["graph_client"] - else: - url = context.get("url") or "http://localhost:8080" - graph = context.get("graph") or "hugegraph" - user = context.get("user") or "admin" - pwd = context.get("pwd") or "admin" - gs = context.get("graphspace") or None - self._client = PyHugeClient(url, graph, user, pwd, gs) - assert self._client is not None, "No valid graph to search." - - def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: - query = context["query"] - vertices = context.get("match_vids") - query_embedding = context.get("query_embedding") - - self.operator_list.clear() - self.operator_list.example_index_query( - num_examples=self._num_gremlin_generate_example - ) - gremlin_response = self.operator_list.gremlin_generate_synthesize( - context["simple_schema"], - vertices=vertices, - gremlin_prompt=self.gremlin_prompt, - ).run(query=query, query_embedding=query_embedding) - if self._num_gremlin_generate_example > 0: - gremlin = gremlin_response["result"] - else: - gremlin = gremlin_response["raw_result"] - log.info("Generated gremlin: %s", gremlin) - context["gremlin"] = gremlin - try: - result = self._client.gremlin().exec(gremlin=gremlin)["data"] - if result == [None]: - result = [] - context["graph_result"] = [ - json.dumps(item, ensure_ascii=False) for item in result - ] - if context["graph_result"]: - context["graph_result_flag"] = 1 - context["graph_context_head"] = ( - f"The following are graph query result " - f"from gremlin query `{gremlin}`.\n" - ) - except Exception as e: # pylint: disable=broad-except - log.error(e) - context["graph_result"] = [] - return context - - def _limit_property_query( - self, value: Optional[str], item_type: str - ) -> Optional[str]: - # NOTE: we skip the filter for list/set type (e.g., list of string, add it if needed) - if not self._limit_property or not isinstance(value, str): - return value - - max_len = self._max_v_prop_len if item_type == "v" else self._max_e_prop_len - return value[:max_len] if value else value - - def _process_vertex( - self, - item: Any, - flat_rel: str, - node_cache: Set[str], - prior_edge_str_len: int, - depth: int, - nodes_with_degree: List[str], - use_id_to_match: bool, - v_cache: Set[str], - ) -> Tuple[str, int, int]: - matched_str = ( - item["id"] if use_id_to_match else item["props"][self._prop_to_match] - ) - if matched_str in node_cache: - flat_rel = flat_rel[:-prior_edge_str_len] - return flat_rel, prior_edge_str_len, depth - - node_cache.add(matched_str) - props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'v')}" - for k, v in item["props"].items() - if v - ) - - # TODO: we may remove label id or replace with label name - if matched_str in v_cache: - node_str = matched_str - else: - v_cache.add(matched_str) - node_str = f"{item['id']}{{{props_str}}}" - - flat_rel += node_str - nodes_with_degree.append(node_str) - depth += 1 - return flat_rel, prior_edge_str_len, depth - - def _process_edge( - self, - item: Any, - path_str: str, - raw_flat_rel: List[Any], - i: int, - use_id_to_match: bool, - e_cache: Set[Tuple[str, str, str]], - ) -> Tuple[str, int]: - props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'e')}" - for k, v in item["props"].items() - if v - ) - props_str = f"{{{props_str}}}" if props_str else "" - prev_matched_str = ( - raw_flat_rel[i - 1]["id"] - if use_id_to_match - else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] - ) - - edge_key = (item["inV"], item["label"], item["outV"]) - if edge_key not in e_cache: - e_cache.add(edge_key) - edge_label = f"{item['label']}{props_str}" - else: - edge_label = item["label"] - - edge_str = ( - f"--[{edge_label}]-->" - if item["outV"] == prev_matched_str - else f"<--[{edge_label}]--" - ) - path_str += edge_str - prior_edge_str_len = len(edge_str) - return path_str, prior_edge_str_len - - def _process_path( - self, - path: Any, - use_id_to_match: bool, - v_cache: Set[str], - e_cache: Set[Tuple[str, str, str]], - ) -> Tuple[str, List[str]]: - flat_rel = "" - raw_flat_rel = path["objects"] - assert len(raw_flat_rel) % 2 == 1, "The length of raw_flat_rel should be odd." - - node_cache = set() - prior_edge_str_len = 0 - depth = 0 - nodes_with_degree = [] - - for i, item in enumerate(raw_flat_rel): - if i % 2 == 0: - # Process each vertex - flat_rel, prior_edge_str_len, depth = self._process_vertex( - item, - flat_rel, - node_cache, - prior_edge_str_len, - depth, - nodes_with_degree, - use_id_to_match, - v_cache, - ) - else: - # Process each edge - flat_rel, prior_edge_str_len = self._process_edge( - item, flat_rel, raw_flat_rel, i, use_id_to_match, e_cache - ) - - return flat_rel, nodes_with_degree - - def _update_vertex_degree_list( - self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str] - ) -> None: - for depth, node_str in enumerate(nodes_with_degree): - if depth >= len(vertex_degree_list): - vertex_degree_list.append(set()) - vertex_degree_list[depth].add(node_str) - - def _format_graph_query_result( - self, query_paths - ) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: - use_id_to_match = self._prop_to_match is None - subgraph = set() - subgraph_with_degree = {} - vertex_degree_list: List[Set[str]] = [] - v_cache: Set[str] = set() - e_cache: Set[Tuple[str, str, str]] = set() - - for path in query_paths: - # 1. Process each path - path_str, vertex_with_degree = self._process_path( - path, use_id_to_match, v_cache, e_cache - ) - subgraph.add(path_str) - subgraph_with_degree[path_str] = vertex_with_degree - # 2. Update vertex degree list - self._update_vertex_degree_list(vertex_degree_list, vertex_with_degree) - - return subgraph, vertex_degree_list, subgraph_with_degree - - def _get_graph_schema(self, refresh: bool = False) -> str: - if self._schema and not refresh: - return self._schema - - schema = self._client.schema() - vertex_schema = schema.getVertexLabels() - edge_schema = schema.getEdgeLabels() - relationships = schema.getRelations() - - self._schema = ( - f"Vertex properties: {vertex_schema}\n" - f"Edge properties: {edge_schema}\n" - f"Relationships: {relationships}\n" - ) - log.debug("Link(Relation): %s", relationships) - return self._schema - - @staticmethod - def _extract_label_names( - source: str, head: str = "name: ", tail: str = ", " - ) -> List[str]: - result = [] - for s in source.split(head): - end = s.find(tail) - label = s[:end] - if label: - result.append(label) - return result - - def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: - schema = self._get_graph_schema() - vertex_props_str, edge_props_str = schema.split("\n")[:2] - # TODO: rename to vertex (also need update in the schema) - vertex_props_str = ( - vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") - ) - edge_props_str = ( - edge_props_str[len("Edge properties: ") :].strip("[").strip("]") - ) - vertex_labels = self._extract_label_names(vertex_props_str) - edge_labels = self._extract_label_names(edge_props_str) - return vertex_labels, edge_labels - - def _format_graph_from_vertex(self, query_result: List[Any]) -> Set[str]: - knowledge = set() - for item in query_result: - props_str = ", ".join(f"{k}: {v}" for k, v in item["properties"].items()) - node_str = f"{item['id']}{{{props_str}}}" - knowledge.add(node_str) - return knowledge - - def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: - # 1. Extract params from context - matched_vids = context.get("match_vids") - if isinstance(context.get("max_deep"), int): - self._max_deep = context["max_deep"] - if isinstance(context.get("max_items"), int): - self._max_items = context["max_items"] - if isinstance(context.get("prop_to_match"), str): - self._prop_to_match = context["prop_to_match"] - - # 2. Extract edge_labels from graph schema - _, edge_labels = self._extract_labels_from_schema() - edge_labels_str = ",".join("'" + label + "'" for label in edge_labels) - # TODO: enhance the limit logic later - edge_limit_amount = len(edge_labels) * huge_settings.edge_limit_pre_label - - use_id_to_match = self._prop_to_match is None - if use_id_to_match: - if not matched_vids: - return context - - gremlin_query = VERTEX_QUERY_TPL.format(keywords=matched_vids) - vertexes = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - log.debug("Vids gremlin query: %s", gremlin_query) - - vertex_knowledge = self._format_graph_from_vertex(query_result=vertexes) - paths: List[Any] = [] - # TODO: use generator or asyncio to speed up the query logic - for matched_vid in matched_vids: - gremlin_query = VID_QUERY_NEIGHBOR_TPL.format( - keywords=f"'{matched_vid}'", - max_deep=self._max_deep, - edge_labels=edge_labels_str, - edge_limit=edge_limit_amount, - max_items=self._max_items, - ) - log.debug("Kneighbor gremlin query: %s", gremlin_query) - paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) - - ( - graph_chain_knowledge, - vertex_degree_list, - knowledge_with_degree, - ) = self._format_graph_query_result(query_paths=paths) - - # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) - if not graph_chain_knowledge: - graph_chain_knowledge.update(vertex_knowledge) - if vertex_degree_list: - vertex_degree_list[0].update(vertex_knowledge) - else: - vertex_degree_list.append(vertex_knowledge) - else: - # WARN: When will the query enter here? - keywords = context.get("keywords") - assert keywords, "No related property(keywords) for graph query." - keywords_str = ",".join("'" + kw + "'" for kw in keywords) - gremlin_query = PROPERTY_QUERY_NEIGHBOR_TPL.format( - prop=self._prop_to_match, - keywords=keywords_str, - edge_labels=edge_labels_str, - edge_limit=edge_limit_amount, - max_deep=self._max_deep, - max_items=self._max_items, - ) - log.warning( - "Unable to find vid, downgraded to property query, please confirm if it meets expectation." - ) - - paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)[ - "data" - ] - ( - graph_chain_knowledge, - vertex_degree_list, - knowledge_with_degree, - ) = self._format_graph_query_result(query_paths=paths) - - context["graph_result"] = list(graph_chain_knowledge) - if context["graph_result"]: - context["graph_result_flag"] = 0 - context["vertex_degree_list"] = [ - list(vertex_degree) for vertex_degree in vertex_degree_list - ] - context["knowledge_with_degree"] = knowledge_with_degree - context["graph_context_head"] = ( - f"The following are graph knowledge in {self._max_deep} depth, e.g:\n" - "`vertexA--[links]-->vertexB<--[links]--vertexC ...`" - "extracted based on key entities as subject:\n" - ) - return context - def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: """ Execute the graph query operation. @@ -471,30 +76,16 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: return data_json # Execute the graph query (assuming schema and semantic query have been completed in previous nodes) - self.init_client(data_json) - - # initial flag: -1 means no result, 0 means subgraph query, 1 means gremlin query - data_json["graph_result_flag"] = -1 - # 1. Try to perform a query based on the generated gremlin - if self._num_gremlin_generate_example >= 0: - data_json = self._gremlin_generate_query(data_json) - # 2. Try to perform a query based on subgraph-search if the previous query failed - if not data_json.get("graph_result"): - data_json = self._subgraph_query(data_json) - - if data_json.get("graph_result"): - log.debug( - "Knowledge from Graph:\n%s", "\n".join(data_json["graph_result"]) - ) - else: - log.debug("No Knowledge Extracted from Graph") + graph_result = self.graph_rag_query.run(data_json) + data_json.update(graph_result) log.info( - f"Graph query completed, found {len(data_json.get('graph_result', []))} results" + "Graph query completed, found %d results", + len(data_json.get("graph_result", [])), ) return data_json except Exception as e: - log.error(f"Graph query failed: {e}") + log.error("Graph query failed: %s", e) return 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 c01cffc91..6a4424f38 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 @@ -13,8 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from hugegraph_llm.config import llm_settings -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -26,7 +25,7 @@ class BuildSemanticIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): - self.build_semantic_index_op = BuildSemanticIndex(get_embedding(llm_settings)) + self.build_semantic_index_op = BuildSemanticIndex(Embeddings().get_embedding()) return super().node_init() def operator_schedule(self, data_json): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py index 1f6a3c75b..28f2cb041 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py @@ -13,8 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from hugegraph_llm.config import llm_settings -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -26,7 +25,7 @@ class BuildVectorIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): - self.build_vector_index_op = BuildVectorIndex(get_embedding(llm_settings)) + self.build_vector_index_op = BuildVectorIndex(Embeddings().get_embedding()) return super().node_init() def operator_schedule(self, data_json): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py index e9283598a..8b9e0db4d 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -19,12 +19,11 @@ from PyCGraph import CStatus -from hugegraph_llm.config import llm_settings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.gremlin_example_index_query import ( GremlinExampleIndexQuery, ) -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings class GremlinExampleIndexQueryNode(BaseNode): @@ -32,15 +31,13 @@ class GremlinExampleIndexQueryNode(BaseNode): def node_init(self): # Build operator (index lazy-loading handled in operator) - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() example_num = getattr(self.wk_input, "example_num", None) if not isinstance(example_num, int): example_num = 2 # Clamp to [0, 10] example_num = max(0, min(10, example_num)) - self.operator = GremlinExampleIndexQuery( - embedding=embedding, num_examples=example_num - ) + self.operator = GremlinExampleIndexQuery(embedding=embedding, num_examples=example_num) return CStatus() def operator_schedule(self, data_json: Dict[str, Any]): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py index 18e480b12..ba5261c59 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -13,12 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from typing import Dict, Any + +from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.semantic_id_query import SemanticIdQuery -from hugegraph_llm.models.embeddings.init_embedding import get_embedding -from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.models.embeddings.init_embedding import Embeddings +from hugegraph_llm.config import huge_settings from hugegraph_llm.utils.log import log @@ -38,22 +39,10 @@ def node_init(self): if not graph_name: return CStatus(-1, "graph_name is required in wk_input") - embedding = get_embedding(llm_settings) - by = ( - self.wk_input.semantic_by - if self.wk_input.semantic_by is not None - else "keywords" - ) - topk_per_keyword = ( - self.wk_input.topk_per_keyword - if self.wk_input.topk_per_keyword is not None - else huge_settings.topk_per_keyword - ) - topk_per_query = ( - self.wk_input.topk_per_query - if self.wk_input.topk_per_query is not None - else 10 - ) + embedding = Embeddings().get_embedding() + by = self.wk_input.semantic_by or "keywords" + topk_per_keyword = self.wk_input.topk_per_keyword or huge_settings.topk_per_keyword + topk_per_query = self.wk_input.topk_per_query or 10 vector_dis_threshold = ( self.wk_input.vector_dis_threshold if self.wk_input.vector_dis_threshold is not None @@ -71,7 +60,7 @@ def node_init(self): return super().node_init() except Exception as e: - log.error(f"Failed to initialize SemanticIdQueryNode: {e}") + log.error("Failed to initialize SemanticIdQueryNode: %s", e) return CStatus(-1, f"SemanticIdQueryNode initialization failed: {e}") @@ -92,12 +81,10 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: semantic_result = self.semantic_id_query.run(data_json) match_vids = semantic_result.get("match_vids", []) - log.info( - f"Semantic query completed, found {len(match_vids)} matching vertex IDs" - ) + log.info("Semantic query completed, found %d matching vertex IDs", len(match_vids)) return semantic_result except Exception as e: - log.error(f"Semantic query failed: {e}") + log.error("Semantic query failed: %s", e) return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py index 48b50acf3..f9af6c49d 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py @@ -14,10 +14,9 @@ # limitations under the License. from typing import Dict, Any -from hugegraph_llm.config import llm_settings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.utils.log import log @@ -34,15 +33,13 @@ def node_init(self): """ try: # 从 wk_input 中读取用户配置参数 - embedding = get_embedding(llm_settings) - max_items = ( - self.wk_input.max_items if self.wk_input.max_items is not None else 3 - ) + embedding = Embeddings().get_embedding() + max_items = self.wk_input.max_items if self.wk_input.max_items is not None else 3 self.operator = VectorIndexQuery(embedding=embedding, topk=max_items) return super().node_init() except Exception as e: - log.error(f"Failed to initialize VectorQueryNode: {e}") + log.error("Failed to initialize VectorQueryNode: %s", e) from PyCGraph import CStatus return CStatus(-1, f"VectorQueryNode initialization failed: {e}") @@ -70,5 +67,5 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: return data_json except Exception as e: - log.error(f"Vector query failed: {e}") + log.error("Vector query failed: %s", e) return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py index 22b970b4a..cc3a7c25a 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py @@ -46,7 +46,7 @@ def node_init(self): ) return super().node_init() except Exception as e: - log.error(f"Failed to initialize AnswerSynthesizeNode: {e}") + log.error("Failed to initialize AnswerSynthesizeNode: %s", e) from PyCGraph import CStatus return CStatus(-1, f"AnswerSynthesizeNode initialization failed: {e}") @@ -75,9 +75,7 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: if result.get("graph_vector_answer"): answer_types.append("graph_vector") - log.info( - f"Answer synthesis completed for types: {', '.join(answer_types)}" - ) + log.info("Answer synthesis completed for types: %s", ", ".join(answer_types)) # Print enabled answer types according to self.wk_input configuration wk_input_types = [] @@ -95,5 +93,5 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: return result except Exception as e: - log.error(f"Answer synthesis failed: {e}") + log.error("Answer synthesis failed: %s", e) return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py index 76fc06eb3..cb390b7b7 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py @@ -34,15 +34,9 @@ def node_init(self): """ try: max_keywords = ( - self.wk_input.max_keywords - if self.wk_input.max_keywords is not None - else 5 - ) - language = ( - self.wk_input.language - if self.wk_input.language is not None - else "english" + self.wk_input.max_keywords if self.wk_input.max_keywords is not None else 5 ) + language = self.wk_input.language if self.wk_input.language is not None else "english" extract_template = self.wk_input.keywords_extract_prompt self.operator = KeywordExtract( @@ -53,7 +47,7 @@ def node_init(self): ) return super().node_init() except Exception as e: - log.error(f"Failed to initialize KeywordExtractNode: {e}") + log.error("Failed to initialize KeywordExtractNode: %s", e) return CStatus(-1, f"KeywordExtractNode initialization failed: {e}") def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: @@ -67,12 +61,12 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: log.warning("Keyword extraction result missing 'keywords' field") result["keywords"] = [] - log.info(f"Extracted keywords: {result.get('keywords', [])}") + log.info("Extracted keywords: %s", result.get("keywords", [])) return result except Exception as e: - log.error(f"Keyword extraction failed: {e}") + log.error("Keyword extraction failed: %s", e) # Add error flag to indicate failure error_result = data_json.copy() error_result["error"] = str(e) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py index 0904b9920..5ea90b570 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py @@ -26,14 +26,11 @@ def _stable_schema_string(state_json: Dict[str, Any]) -> str: - val = state_json.get("simple_schema") - if val is None: - val = state_json.get("schema") - if val is None: - return "" - if isinstance(val, str): - return val - return json.dumps(val, ensure_ascii=False, sort_keys=True) + if "simple_schema" in state_json and state_json["simple_schema"] is not None: + return json.dumps(state_json["simple_schema"], ensure_ascii=False, sort_keys=True) + if "schema" in state_json and state_json["schema"] is not None: + return json.dumps(state_json["schema"], ensure_ascii=False, sort_keys=True) + return "" class Text2GremlinNode(BaseNode): diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index fc729c11e..47b0f060f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -16,7 +16,7 @@ # under the License. -from typing import Any, Optional, Dict +from typing import Any, Dict, Optional from hugegraph_llm.enums.property_cardinality import PropertyCardinality from hugegraph_llm.enums.property_data_type import PropertyDataType diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py index e93d916b3..4c4c167c4 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/fetch_graph_data.py @@ -22,7 +22,6 @@ class FetchGraphData: - def __init__(self, graph: PyHugeClient): self.graph = graph 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 90f1c00ea..2f0643a77 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 @@ -33,7 +33,7 @@ def __init__(self, graph_name: str): self.schema = self.client.schema() def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: - mini_schema = {} + mini_schema = {} # type: ignore # Add necessary vertexlabels items (3) if "vertexlabels" in schema: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index 6d9f96214..fb385a59c 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -17,31 +17,25 @@ import asyncio -import os -from typing import Dict, Any, List +from typing import Any, Dict, List -from hugegraph_llm.config import resource_path, llm_settings, huge_settings -from hugegraph_llm.indices.vector_index import VectorIndex +from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import ( - get_embeddings_parallel, - get_filename_prefix, - get_index_folder_name, -) +from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel # FIXME: we need keep the logic same with build_semantic_index.py class BuildGremlinExampleIndex: - def __init__(self, embedding: BaseEmbedding, examples: List[Dict[str, str]]): - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "gremlin_examples")) + def __init__( + self, + embedding: BaseEmbedding, + examples: List[Dict[str, str]], + vector_index: type[VectorStoreBase], + ): + self.vector_index_name = "gremlin_examples" self.examples = examples self.embedding = embedding - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(embedding, "model_name", None) - ) + self.vector_index = vector_index def run(self, context: Dict[str, Any]) -> Dict[str, Any]: # !: We have assumed that self.example is not empty @@ -50,8 +44,8 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: examples_embedding = asyncio.run(get_embeddings_parallel(self.embedding, queries)) embed_dim = len(examples_embedding[0]) if len(self.examples) > 0: - vector_index = VectorIndex(embed_dim) + vector_index = self.vector_index.from_name(embed_dim, self.vector_index_name) vector_index.add(examples_embedding, self.examples) - vector_index.to_index_file(self.index_dir, self.filename_prefix) + vector_index.save_index_by_name(self.vector_index_name) context["embed_dim"] = embed_dim return context 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 2ed4e840a..82b8b325f 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 @@ -15,36 +15,22 @@ # specific language governing permissions and limitations # under the License. - import asyncio -import os from typing import Any, Dict -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index import VectorIndex +from tqdm import tqdm + +from hugegraph_llm.config import huge_settings +from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager -from hugegraph_llm.utils.embedding_utils import ( - get_embeddings_parallel, - get_filename_prefix, - get_index_folder_name, -) from hugegraph_llm.utils.log import log class BuildSemanticIndex: - def __init__(self, embedding: BaseEmbedding): - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str( - os.path.join(resource_path, self.folder_name, "graph_vids") - ) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(embedding, "model_name", None) - ) - self.vid_index = VectorIndex.from_index_file( - self.index_dir, self.filename_prefix + 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" ) self.embedding = embedding self.sm = SchemaManager(huge_settings.graph_name) @@ -52,32 +38,47 @@ def __init__(self, embedding: BaseEmbedding): def _extract_names(self, vertices: list[str]) -> list[str]: return [v.split(":")[1] for v in vertices] + async def _get_embeddings_parallel(self, vids: list[str]) -> list[Any]: + sem = asyncio.Semaphore(10) + batch_size = 1000 + + async def get_embeddings_with_semaphore(vid_list: list[str]) -> Any: + async with sem: + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, self.embedding.get_texts_embeddings, vid_list + ) + + vid_batches = [vids[i : i + batch_size] for i in range(0, len(vids), batch_size)] + tasks = [get_embeddings_with_semaphore(batch) for batch in vid_batches] + + embeddings = [] + with tqdm(total=len(tasks)) as pbar: + for future in asyncio.as_completed(tasks): + batch_embeddings = await future + embeddings.extend(batch_embeddings) + pbar.update(1) + return embeddings + def run(self, context: Dict[str, Any]) -> Dict[str, Any]: vertexlabels = self.sm.schema.getSchema()["vertexlabels"] all_pk_flag = bool(vertexlabels) and all( data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels ) - past_vids = self.vid_index.properties + past_vids = self.vid_index.get_all_properties() # only support Faiss # TODO: We should build vid vector index separately, especially when the vertices may be very large - - present_vids = context[ - "vertices" - ] # Warning: data truncated by fetch_graph_data.py + present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py removed_vids = set(past_vids) - set(present_vids) removed_num = self.vid_index.remove(removed_vids) added_vids = list(set(present_vids) - set(past_vids)) if added_vids: - vids_to_process = ( - self._extract_names(added_vids) if all_pk_flag else added_vids - ) - added_embeddings = asyncio.run( - get_embeddings_parallel(self.embedding, vids_to_process) - ) + vids_to_process = self._extract_names(added_vids) if all_pk_flag else added_vids + 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.to_index_file(self.index_dir, self.filename_prefix) + self.vid_index.save_index_by_name(huge_settings.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/index_op/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py index f5fb823c5..64967d2df 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py @@ -15,43 +15,33 @@ # specific language governing permissions and limitations # under the License. - import asyncio -import os -from typing import Dict, Any +from typing import Any, Dict -from hugegraph_llm.config import huge_settings, resource_path, llm_settings -from hugegraph_llm.indices.vector_index import VectorIndex +from hugegraph_llm.config import huge_settings +from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import ( - get_embeddings_parallel, - get_filename_prefix, - get_index_folder_name, -) +from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel from hugegraph_llm.utils.log import log class BuildVectorIndex: - def __init__(self, embedding: BaseEmbedding): + def __init__(self, embedding: BaseEmbedding, vector_index: type[VectorStoreBase]): self.embedding = embedding - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(self.embedding, "model_name", None) + self.vector_index = vector_index.from_name( + embedding.get_embedding_dim(), + huge_settings.graph_name, + "chunks", ) - self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: if "chunks" not in context: raise ValueError("chunks not found in context.") chunks = context["chunks"] - chunks_embedding = [] log.debug("Building vector index for %s chunks...", len(context["chunks"])) - # TODO: use async_get_texts_embedding instead of single sync method - chunks_embedding = asyncio.run(get_embeddings_parallel(self.embedding, chunks)) + # Use async parallel embedding to speed up + chunks_embedding = asyncio.run(get_embeddings_parallel(self.embedding, chunks)) # type: ignore if len(chunks_embedding) > 0: self.vector_index.add(chunks_embedding, chunks) - self.vector_index.to_index_file(self.index_dir, self.filename_prefix) + self.vector_index.save_index_by_name(huge_settings.graph_name, "chunks") return context diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index b680f2ca3..e3eea9f07 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -16,53 +16,38 @@ # under the License. -import asyncio import os -from typing import Dict, Any, List +from typing import Any, Dict, List, Optional import pandas as pd +from tqdm import tqdm -from hugegraph_llm.config import resource_path, llm_settings, huge_settings -from hugegraph_llm.indices.vector_index import VectorIndex, INDEX_FILE_NAME, PROPERTIES_FILE_NAME +from hugegraph_llm.config import resource_path +from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.utils.embedding_utils import ( - get_embeddings_parallel, - get_filename_prefix, - get_index_folder_name, -) from hugegraph_llm.utils.log import log class GremlinExampleIndexQuery: - def __init__(self, embedding: BaseEmbedding = None, num_examples: int = 1): + def __init__( + self, + vector_index: type[VectorStoreBase], + embedding: Optional[BaseEmbedding] = None, + num_examples: int = 1, + ): self.embedding = embedding or Embeddings().get_embedding() self.num_examples = num_examples - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "gremlin_examples")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(self.embedding, "model_name", None) - ) - self._ensure_index_exists() - self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) - - def _ensure_index_exists(self): - index_name = ( - f"{self.filename_prefix}_{INDEX_FILE_NAME}" if self.filename_prefix else INDEX_FILE_NAME - ) - props_name = ( - f"{self.filename_prefix}_{PROPERTIES_FILE_NAME}" - if self.filename_prefix - else PROPERTIES_FILE_NAME - ) - if not ( - os.path.exists(os.path.join(self.index_dir, index_name)) - and os.path.exists(os.path.join(self.index_dir, props_name)) - ): + if not vector_index.exist("gremlin_examples"): log.warning("No gremlin example index found, will generate one.") + self.vector_index = vector_index.from_name( + self.embedding.get_embedding_dim(), "gremlin_examples" + ) self._build_default_example_index() + else: + self.vector_index = vector_index.from_name( + self.embedding.get_embedding_dim(), "gremlin_examples" + ) def _get_match_result(self, context: Dict[str, Any], query: str) -> List[Dict[str, Any]]: if self.num_examples <= 0: @@ -77,12 +62,20 @@ def _build_default_example_index(self): properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict( orient="records" ) + from concurrent.futures import ThreadPoolExecutor + # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) - queries = [row["query"] for row in properties] - embeddings = asyncio.run(get_embeddings_parallel(self.embedding, queries)) - vector_index = VectorIndex(len(embeddings[0])) - vector_index.add(embeddings, properties) - vector_index.to_index_file(self.index_dir, self.filename_prefix) + with ThreadPoolExecutor() as executor: + embeddings = list( + tqdm( + executor.map( + self.embedding.get_text_embedding, [row["query"] for row in properties] + ), + total=len(properties), + ) + ) + self.vector_index.add(embeddings, properties) + self.vector_index.save_index_by_name("gremlin_examples") def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query = context.get("query") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py index 3ac03246f..49e712d01 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py @@ -17,14 +17,14 @@ import os -from typing import Dict, Any, Literal, List, Tuple +from typing import Any, Dict, List, Literal, Tuple -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index import VectorIndex +from pyhugegraph.client import PyHugeClient + +from hugegraph_llm.config import huge_settings, resource_path +from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import get_filename_prefix, get_index_folder_name from hugegraph_llm.utils.log import log -from pyhugegraph.client import PyHugeClient class SemanticIdQuery: @@ -33,19 +33,16 @@ class SemanticIdQuery: def __init__( self, embedding: BaseEmbedding, + vector_index: type[VectorStoreBase], by: Literal["query", "keywords"] = "keywords", topk_per_query: int = 10, topk_per_keyword: int = huge_settings.topk_per_keyword, vector_dis_threshold: float = huge_settings.vector_dis_threshold, ): - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(embedding, "model_name", None) + self.index_dir = str(os.path.join(resource_path, huge_settings.graph_name, "graph_vids")) + self.vector_index = vector_index.from_name( + embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids" ) - self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) self.embedding = embedding self.by = by self.topk_per_query = topk_per_query @@ -82,7 +79,7 @@ def _exact_match_vids(self, keywords: List[str]) -> Tuple[List[str], List[str]]: def _fuzzy_match_vids(self, keywords: List[str]) -> List[str]: fuzzy_match_result = [] for keyword in keywords: - keyword_vector = self.embedding.get_texts_embeddings([keyword])[0] + keyword_vector = self.embedding.get_text_embedding(keyword) results = self.vector_index.search( keyword_vector, top_k=self.topk_per_keyword, @@ -96,7 +93,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: graph_query_list = set() if self.by == "query": query = context["query"] - query_vector = self.embedding.get_texts_embeddings([query])[0] + query_vector = self.embedding.get_text_embedding(query) results = self.vector_index.search(query_vector, top_k=self.topk_per_query) if results: graph_query_list.update(results[: self.topk_per_query]) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py index 4ed616929..5ced65b41 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py @@ -16,28 +16,23 @@ # under the License. -import os -from typing import Dict, Any +from typing import Any, Dict -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index import VectorIndex +from hugegraph_llm.config import huge_settings +from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import get_filename_prefix, get_index_folder_name from hugegraph_llm.utils.log import log class VectorIndexQuery: - def __init__(self, embedding: BaseEmbedding, topk: int = 3): + def __init__( + self, vector_index: type[VectorStoreBase], embedding: BaseEmbedding, topk: int = 3 + ): self.embedding = embedding self.topk = topk - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space + self.vector_index = vector_index.from_name( + embedding.get_embedding_dim(), huge_settings.graph_name, "chunks" ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(embedding, "model_name", None) - ) - self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query = context.get("query") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py index 9138f9e9b..2140abca2 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py @@ -62,8 +62,8 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = ( - f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") + context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( + "\n" ) final_prompt = self._prompt_template.format( @@ -119,8 +119,8 @@ async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[st context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = ( - f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") + context_str = f"{context_head_str}\n{self._context_body}\n{context_tail_str}".strip( + "\n" ) final_prompt = self._prompt_template.format( @@ -133,7 +133,11 @@ async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[st graph_result_context, vector_result_context = self.handle_vector_graph(context) async for context in self.async_streaming_generate( - context, context_head_str, context_tail_str, vector_result_context, graph_result_context + context, + context_head_str, + context_tail_str, + vector_result_context, + graph_result_context, ): yield context @@ -151,10 +155,8 @@ async def async_generate( final_prompt = self._question async_tasks["raw_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) if self._vector_only_answer: - context_str = ( - f"{context_head_str}\n" - f"{vector_result_context}\n" - f"{context_tail_str}".strip("\n") + context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip( + "\n" ) final_prompt = self._prompt_template.format( @@ -164,10 +166,8 @@ async def async_generate( self._llm.agenerate(prompt=final_prompt) ) if self._graph_only_answer: - context_str = ( - f"{context_head_str}\n" - f"{graph_result_context}\n" - f"{context_tail_str}".strip("\n") + context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( + "\n" ) final_prompt = self._prompt_template.format( @@ -180,9 +180,7 @@ async def async_generate( context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = ( - f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") - ) + context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question @@ -235,10 +233,8 @@ async def async_streaming_generate( ) auto_id += 1 if self._vector_only_answer: - context_str = ( - f"{context_head_str}\n" - f"{vector_result_context}\n" - f"{context_tail_str}".strip("\n") + context_str = f"{context_head_str}\n{vector_result_context}\n{context_tail_str}".strip( + "\n" ) final_prompt = self._prompt_template.format( @@ -246,15 +242,15 @@ async def async_streaming_generate( ) async_generators.append( self.__llm_generate_with_meta_info( - task_id=auto_id, target_key="vector_only_answer", prompt=final_prompt + task_id=auto_id, + target_key="vector_only_answer", + prompt=final_prompt, ) ) auto_id += 1 if self._graph_only_answer: - context_str = ( - f"{context_head_str}\n" - f"{graph_result_context}\n" - f"{context_tail_str}".strip("\n") + context_str = f"{context_head_str}\n{graph_result_context}\n{context_tail_str}".strip( + "\n" ) final_prompt = self._prompt_template.format( @@ -270,16 +266,16 @@ async def async_streaming_generate( context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = ( - f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") - ) + context_str = f"{context_head_str}\n{context_body_str}\n{context_tail_str}".strip("\n") final_prompt = self._prompt_template.format( context_str=context_str, query_str=self._question ) async_generators.append( self.__llm_generate_with_meta_info( - task_id=auto_id, target_key="graph_vector_answer", prompt=final_prompt + task_id=auto_id, + target_key="graph_vector_answer", + prompt=final_prompt, ) ) auto_id += 1 diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py index 2ac2eafff..819ef25fe 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py @@ -52,8 +52,7 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: data["triples"] = [] extract_triples_by_regex(llm_output, data) print( - f"LLM {self.__class__.__name__} input:{prompt} \n" - f" output: {llm_output} \n data: {data}" + f"LLM {self.__class__.__name__} input:{prompt} \n output: {llm_output} \n data: {data}" ) data["call_count"] = data.get("call_count", 0) + 1 diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index 650834300..39f49dbdd 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -18,7 +18,7 @@ import asyncio import json import re -from typing import Optional, List, Dict, Any, Union +from typing import Any, Dict, List, Optional, Union from hugegraph_llm.config import prompt from hugegraph_llm.models.llms.base import BaseLLM @@ -29,7 +29,7 @@ class GremlinGenerateSynthesize: def __init__( self, - llm: BaseLLM = None, + llm: BaseLLM | None = None, schema: Optional[Union[dict, str]] = None, vertices: Optional[List[str]] = None, gremlin_prompt: Optional[str] = None, @@ -54,8 +54,7 @@ def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional example_strings = [] for example in examples: example_strings.append( - f"- query: {example['query']}\n" - f"- gremlin:\n```gremlin\n{example['gremlin']}\n```" + f"- query: {example['query']}\n- gremlin:\n```gremlin\n{example['gremlin']}\n```" ) return "\n\n".join(example_strings) 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 8897e0fea..a9ac4c050 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 @@ -16,7 +16,7 @@ # under the License. import re -from typing import List, Any, Dict, Optional +from typing import Any, Dict, List, Optional from hugegraph_llm.document.chunk_split import ChunkSplitter from hugegraph_llm.models.llms.base import BaseLLM @@ -27,11 +27,11 @@ ## Basic Rules 1. The output format must be: (X,Y,Z) - LABEL -In this format, Y must be a value from "properties" or "edge_label", +In this format, Y must be a value from "properties" or "edge_label", and LABEL must be X's vertex_label or Y's edge_label. 2. Don't extract attribute/property fields that do not exist in the given schema 3. Ensure the extract property is in the same type as the schema (like 'age' should be a number) -4. Translate the given schema filed into Chinese if the given text is Chinese but the schema is in English (Optional) +4. Translate the given schema filed into Chinese if the given text is Chinese but the schema is in English (Optional) ## Example (Note: Update the example to correspond to the given text and schema) ### Input example: @@ -76,8 +76,7 @@ def generate_extract_triple_prompt(text, schema=None) -> str: if schema: return schema_real_prompt log.warning( - "Recommend to provide a graph schema to improve the extraction accuracy. " - "Now using the default schema." + "Recommend to provide a graph schema to improve the extraction accuracy. Now using the default schema." ) return text_based_prompt @@ -150,7 +149,7 @@ def extract_triples_by_regex_with_schema(schema, text, graph): } ) break - graph["vertices"] = list(vertices_dict.values()) + graph["vertices"] = vertices_dict.values() class InfoExtract: 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 565d79023..f4933d676 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 @@ -125,8 +125,7 @@ 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 object found, " - "please check the output format example in prompt." + "Invalid property graph! No JSON object found, please check the output format example in prompt." ) return [] json_str = json_match.group(1).strip() diff --git a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py index 129c36fc5..702258ace 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py @@ -16,6 +16,7 @@ # under the License. from typing import Optional, List, Literal, Union +from hugegraph_llm.config import huge_settings from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.operators.common_op.check_schema import CheckSchema @@ -43,7 +44,6 @@ from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery from hugegraph_llm.operators.common_op.merge_dedup_rerank import MergeDedupRerank from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize -from hugegraph_llm.config import huge_settings from pyhugegraph.client import PyHugeClient diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 9c59a42e6..ffd313a73 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -13,9 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import GParam, CStatus +from typing import Union, List, Optional, Any -from typing import Union, List, Optional, Any, Dict +from PyCGraph import GParam, CStatus class WkFlowInput(GParam): @@ -78,7 +78,7 @@ class WkFlowInput(GParam): is_vector_only: bool = False # used for build text2gremin index - examples: Optional[List[Dict[str, str]]] = None + examples: Optional[List[dict]] = None def reset(self, _: CStatus) -> None: self.texts = None @@ -206,11 +206,7 @@ def to_json(self): dict: A dictionary containing non-None instance members and their serialized values. """ # Only export instance attributes (excluding methods and class attributes) whose values are not None - return { - k: v - for k, v in self.__dict__.items() - if not k.startswith("_") and v is not None - } + return {k: v for k, v in self.__dict__.items() if not k.startswith("_") and v is not None} # Implement a method that assigns keys from data_json as WkFlowState member variables def assign_from_json(self, data_json: dict): diff --git a/hugegraph-llm/src/hugegraph_llm/utils/anchor.py b/hugegraph-llm/src/hugegraph_llm/utils/anchor.py index 4542a7fd9..a46a4e499 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/anchor.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/anchor.py @@ -35,6 +35,5 @@ def get_project_root() -> Path: return parent # Raise an error if no project root is found raise RuntimeError( - "Project root could not be determined. " - "Ensure that 'pyproject.toml' or '.git' exists in the project directory." + "Project root could not be determined. Ensure that 'pyproject.toml' or '.git' exists in the project directory." ) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/log.py b/hugegraph-llm/src/hugegraph_llm/utils/log.py old mode 100755 new mode 100644 diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 301a6bdab..9468f823b 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -16,19 +16,18 @@ # under the License. import json -import os +from typing import Type import docx import gradio as gr -from hugegraph_llm.config import resource_path, huge_settings, llm_settings -from hugegraph_llm.indices.vector_index import VectorIndex -from hugegraph_llm.models.embeddings.init_embedding import model_map -from hugegraph_llm.flows.scheduler import SchedulerSingleton -from hugegraph_llm.utils.embedding_utils import ( - get_filename_prefix, - get_index_folder_name, -) +from hugegraph_llm.config import huge_settings, index_settings +from hugegraph_llm.indices.vector_index.base import VectorStoreBase +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex +from hugegraph_llm.models.embeddings.init_embedding import Embeddings +from hugegraph_llm.models.llms.init_llm import LLMs +from hugegraph_llm.operators.kg_construction_task import KgBuilder +from hugegraph_llm.utils.hugegraph_utils import get_hg_client def read_documents(input_file, input_text): @@ -60,26 +59,15 @@ def read_documents(input_file, input_text): # pylint: disable=C0301 def get_vector_index_info(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) - ) - chunk_vector_index = VectorIndex.from_index_file( - str(os.path.join(resource_path, folder_name, "chunks")), - filename_prefix, - record_miss=False, - ) - graph_vid_vector_index = VectorIndex.from_index_file( - str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix + vector_index = get_vector_index_class(index_settings.cur_vector_index) + vector_index_entity = vector_index.from_name( + Embeddings().get_embedding().get_embedding_dim(), huge_settings.graph_name, "chunks" ) + return json.dumps( { - "embed_dim": chunk_vector_index.index.d, - "vector_info": { - "chunk_vector_num": chunk_vector_index.index.ntotal, - "graph_vid_vector_num": graph_vid_vector_index.index.ntotal, - "graph_properties_vector_num": len(chunk_vector_index.properties), - }, + **vector_index_entity.get_vector_index_info(), + "cur_vector_index": index_settings.cur_vector_index, }, ensure_ascii=False, indent=2, @@ -87,17 +75,49 @@ def get_vector_index_info(): def clean_vector_index(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) - ) - VectorIndex.clean(str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix) + vector_index = get_vector_index_class(index_settings.cur_vector_index) + vector_index.clean(huge_settings.graph_name, "chunks") gr.Info("Clean vector index successfully!") def build_vector_index(input_file, input_text): + vector_index = get_vector_index_class(index_settings.cur_vector_index) if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) - scheduler = SchedulerSingleton.get_instance() - return scheduler.schedule_flow("build_vector_index", texts) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + context = builder.chunk_split(texts, "paragraph", "zh").build_vector_index(vector_index).run() + return json.dumps(context, ensure_ascii=False, indent=2) + + +def get_vector_index_class(vector_index_str: str) -> Type[VectorStoreBase]: + if vector_index_str == "Faiss": + return FaissVectorIndex # type: ignore[return-value] + if vector_index_str == "Milvus": + try: + from hugegraph_llm.indices.vector_index.milvus_vector_store import ( # pylint: disable=import-outside-toplevel + MilvusVectorIndex, + ) + + return MilvusVectorIndex # type: ignore[return-value] + except Exception as e: # pylint: disable=broad-except + raise gr.Error( + f"Milvus engine selected but dependency not available: {e}.\n" + "Fix it by running: 'uv sync --extra vectordb' (recommended) or install 'pymilvus' manually.\n" + "Alternatively, switch vector engine to Faiss/Qdrant in the UI." + ) + if vector_index_str == "Qdrant": + try: + from hugegraph_llm.indices.vector_index.qdrant_vector_store import ( # pylint: disable=import-outside-toplevel + QdrantVectorIndex, + ) + + return QdrantVectorIndex # type: ignore[return-value] + except Exception as e: # pylint: disable=broad-except + raise gr.Error( + f"Qdrant engine selected but dependency not available: {e}.\n" + "Fix it by running: 'uv sync --extra vectordb' (recommended) or install 'qdrant-client' manually.\n" + "Alternatively, switch vector engine to Faiss/Milvus in the UI." + ) + # Fallback to Faiss + return FaissVectorIndex # type: ignore[return-value] diff --git a/hugegraph-llm/src/tests/indices/test_vector_index.py b/hugegraph-llm/src/tests/indices/test_faiss_vector_index.py similarity index 81% rename from hugegraph-llm/src/tests/indices/test_vector_index.py rename to hugegraph-llm/src/tests/indices/test_faiss_vector_index.py index 0f8fd5f48..fd1eb2a15 100644 --- a/hugegraph-llm/src/tests/indices/test_vector_index.py +++ b/hugegraph-llm/src/tests/indices/test_faiss_vector_index.py @@ -19,16 +19,20 @@ import unittest from pprint import pprint -from hugegraph_llm.indices.vector_index import VectorIndex +from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding class TestVectorIndex(unittest.TestCase): def test_vector_index(self): embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") - data = ["腾讯的合伙人有字节跳动", "谷歌和微软是竞争关系", "美团的合伙人有字节跳动"] + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] data_embedding = [embedder.get_text_embedding(d) for d in data] - index = VectorIndex(1024) + index = FaissVectorIndex(1024) index.add(data_embedding, data) query = "腾讯的合伙人有哪些？" query_vector = embedder.get_text_embedding(query) diff --git a/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py new file mode 100644 index 000000000..b1ac0f209 --- /dev/null +++ b/hugegraph-llm/src/tests/indices/test_milvus_vector_index.py @@ -0,0 +1,100 @@ +# 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 unittest +from pprint import pprint + +from hugegraph_llm.indices.vector_index.milvus_vector_store import MilvusVectorIndex +from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding + +test_name = "test" + + +class TestMilvusVectorIndex(unittest.TestCase): + def tearDown(self): + MilvusVectorIndex.clean(test_name) + + def test_vector_index(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = MilvusVectorIndex.from_name(1024, test_name) + index.add(data_embedding, data) + + query = "腾讯的合伙人有哪些？" + query_vector = embedder.get_text_embedding(query) + results = index.search(query_vector, 2, dis_threshold=1000) + pprint(results) + + self.assertIsNotNone(results) + self.assertLessEqual(len(results), 2) + + def test_save_and_load(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = MilvusVectorIndex.from_name(1024, test_name) + index.add(data_embedding, data) + + index.save_index_by_name(test_name) + + loaded_index = MilvusVectorIndex.from_name(1024, test_name) + + query = "腾讯的合伙人有哪些？" + query_vector = embedder.get_text_embedding(query) + results = loaded_index.search(query_vector, 2, dis_threshold=1000) + + self.assertIsNotNone(results) + self.assertLessEqual(len(results), 2) + + def test_remove_entries(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = MilvusVectorIndex.from_name(1024, test_name) + index.add(data_embedding, data) + + query = "合伙人" + query_vector = embedder.get_text_embedding(query) + initial_results = index.search(query_vector, 3, dis_threshold=1000) + initial_count = len(initial_results) + + remove_count = index.remove(["谷歌和微软是竞争关系"]) + + self.assertEqual(remove_count, 1) + + after_results = index.search(query_vector, 3, dis_threshold=1000) + self.assertLessEqual(len(after_results), initial_count - 1) + self.assertNotIn("谷歌和微软是竞争关系", after_results) diff --git a/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py new file mode 100644 index 000000000..1e0768051 --- /dev/null +++ b/hugegraph-llm/src/tests/indices/test_qdrant_vector_index.py @@ -0,0 +1,102 @@ +# 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 unittest +from pprint import pprint + +from hugegraph_llm.indices.vector_index.qdrant_vector_store import QdrantVectorIndex +from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding + + +class TestQdrantVectorIndex(unittest.TestCase): + def setUp(self): + self.name = "test" + + def tearDown(self): + QdrantVectorIndex.clean(self.name) + + def test_vector_index(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = QdrantVectorIndex.from_name(1024, self.name) + index.add(data_embedding, data) + + query = "腾讯的合伙人有哪些？" + query_vector = embedder.get_text_embedding(query) + results = index.search(query_vector, 2, dis_threshold=100) + pprint(results) + + self.assertIsNotNone(results) + self.assertLessEqual(len(results), 2) + + def test_save_and_load(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = QdrantVectorIndex.from_name(1024, self.name) + index.add(data_embedding, data) + + index.save_index_by_name(self.name) + + loaded_index = QdrantVectorIndex.from_name(1024, self.name) + + query = "腾讯的合伙人有哪些？" + query_vector = embedder.get_text_embedding(query) + results = loaded_index.search(query_vector, 2, dis_threshold=100) + + self.assertIsNotNone(results) + self.assertLessEqual(len(results), 2) + + def test_remove_entries(self): + embedder = OllamaEmbedding("quentinz/bge-large-zh-v1.5") + + data = [ + "腾讯的合伙人有字节跳动", + "谷歌和微软是竞争关系", + "美团的合伙人有字节跳动", + ] + data_embedding = [embedder.get_text_embedding(d) for d in data] + + index = QdrantVectorIndex.from_name(1024, self.name) + index.add(data_embedding, data) + + query = "合伙人" + query_vector = embedder.get_text_embedding(query) + initial_results = index.search(query_vector, 3, dis_threshold=100) + initial_count = len(initial_results) + + remove_count = index.remove(["谷歌和微软是竞争关系"]) + + self.assertEqual(remove_count, 1) + + after_results = index.search(query_vector, 3) + self.assertLessEqual(len(after_results), initial_count - 1) + self.assertNotIn("谷歌和微软是竞争关系", after_results) diff --git a/hugegraph-ml/src/hugegraph_ml/models/bgrl.py b/hugegraph-ml/src/hugegraph_ml/models/bgrl.py index 288a434b8..0000e546c 100644 --- a/hugegraph-ml/src/hugegraph_ml/models/bgrl.py +++ b/hugegraph-ml/src/hugegraph_ml/models/bgrl.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -# pylint: disable=C0103,R1705.R1734 +# pylint: disable=C0103,R1705.R1734,E1102 """ Bootstrapped Graph Latents (BGRL) diff --git a/hugegraph-python-client/src/pyhugegraph/api/auth.py b/hugegraph-python-client/src/pyhugegraph/api/auth.py index ab7d66169..d127c4f6d 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/auth.py +++ b/hugegraph-python-client/src/pyhugegraph/api/auth.py @@ -119,8 +119,8 @@ def revoke_accesses(self, access_id) -> Optional[Dict]: # pylint: disable=unuse @router.http("PUT", "auth/accesses/{access_id}") def modify_accesses( - self, access_id, access_description # pylint: disable=unused-argument - ) -> Optional[Dict]: + self, access_id, access_description + ) -> Optional[Dict]: # pylint: disable=unused-argument # The permission of access can\'t be updated data = {"access_description": access_description} return self._invoke_request(data=json.dumps(data)) @@ -174,8 +174,8 @@ def update_target( @router.http("GET", "auth/targets/{target_id}") def get_target( - self, target_id, response=None # pylint: disable=unused-argument - ) -> Optional[Dict]: + self, target_id, response=None + ) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "auth/targets") @@ -193,8 +193,8 @@ def delete_belong(self, belong_id) -> None: # pylint: disable=unused-argument @router.http("PUT", "auth/belongs/{belong_id}") def update_belong( - self, belong_id, description # pylint: disable=unused-argument - ) -> Optional[Dict]: + self, belong_id, description + ) -> Optional[Dict]: # pylint: disable=unused-argument data = {"belong_description": description} return self._invoke_request(data=json.dumps(data)) diff --git a/hugegraph-python-client/src/pyhugegraph/api/graph.py b/hugegraph-python-client/src/pyhugegraph/api/graph.py index 4555eeda4..2372b6522 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/graph.py +++ b/hugegraph-python-client/src/pyhugegraph/api/graph.py @@ -138,17 +138,13 @@ def addEdges(self, input_data) -> Optional[List[EdgeData]]: return None @router.http("PUT", "graph/edges/{edge_id}?action=append") - def appendEdge( - self, edge_id, properties # pylint: disable=unused-argument - ) -> Optional[EdgeData]: + def appendEdge(self, edge_id, properties): # pylint: disable=unused-argument if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None @router.http("PUT", "graph/edges/{edge_id}?action=eliminate") - def eliminateEdge( - self, edge_id, properties # pylint: disable=unused-argument - ) -> Optional[EdgeData]: + def eliminateEdge(self, edge_id, properties): # pylint: disable=unused-argument if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None diff --git a/hugegraph-python-client/src/pyhugegraph/api/gremlin.py b/hugegraph-python-client/src/pyhugegraph/api/gremlin.py index e02e7fb2a..3261d60b3 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/gremlin.py +++ b/hugegraph-python-client/src/pyhugegraph/api/gremlin.py @@ -43,9 +43,7 @@ def exec(self, gremlin): try: if response := self._invoke_request(data=gremlin_data.to_json()): return ResponseData(response).result - log.error( # pylint: disable=logging-fstring-interpolation - f"Gremlin can't get results: {str(response)}" - ) + log.error("Gremlin can't get results: %s", str(response)) return None except Exception as e: raise NotFoundError(f"Gremlin can't get results: {e}") from e diff --git a/hugegraph-python-client/src/pyhugegraph/api/schema.py b/hugegraph-python-client/src/pyhugegraph/api/schema.py index 7e8926678..8095887b0 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/schema.py +++ b/hugegraph-python-client/src/pyhugegraph/api/schema.py @@ -69,8 +69,8 @@ def getSchema(self, _format: str = "json") -> Optional[Dict]: # pylint: disable @router.http("GET", "schema/propertykeys/{property_name}") def getPropertyKey( - self, property_name # pylint: disable=unused-argument - ) -> Optional[PropertyKeyData]: + self, property_name + ) -> Optional[PropertyKeyData]: # pylint: disable=unused-argument if response := self._invoke_request(): return PropertyKeyData(response) return None @@ -96,8 +96,8 @@ def getVertexLabels(self) -> Optional[List[VertexLabelData]]: @router.http("GET", "schema/edgelabels/{label_name}") def getEdgeLabel( - self, label_name: str # pylint: disable=unused-argument - ) -> Optional[EdgeLabelData]: + self, label_name: str + ) -> Optional[EdgeLabelData]: # pylint: disable=unused-argument if response := self._invoke_request(): return EdgeLabelData(response) log.error("EdgeLabel not found: %s", str(response)) diff --git a/hugegraph-python-client/src/pyhugegraph/api/traverser.py b/hugegraph-python-client/src/pyhugegraph/api/traverser.py index 72dddb07a..2f226522b 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/traverser.py +++ b/hugegraph-python-client/src/pyhugegraph/api/traverser.py @@ -50,8 +50,8 @@ def shortest_path(self, source_id, target_id, max_depth): # pylint: disable=unu 'traversers/allshortestpaths?source="{source_id}"&target="{target_id}"&max_depth={max_depth}', ) def all_shortest_paths( - self, source_id, target_id, max_depth # pylint: disable=unused-argument - ): + self, source_id, target_id, max_depth + ): # pylint: disable=unused-argument return self._invoke_request() @router.http( @@ -60,8 +60,8 @@ def all_shortest_paths( "&weight={weight}&max_depth={max_depth}", ) def weighted_shortest_path( - self, source_id, target_id, weight, max_depth # pylint: disable=unused-argument - ): + self, source_id, target_id, weight, max_depth + ): # pylint: disable=unused-argument return self._invoke_request() @router.http( diff --git a/hugegraph-python-client/src/pyhugegraph/utils/util.py b/hugegraph-python-client/src/pyhugegraph/utils/util.py index 56a135547..76770a818 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/util.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/util.py @@ -58,12 +58,7 @@ def check_if_success(response, error=None): req = response.request req_body = req.body if req.body else "Empty body" response_body = response.text if response.text else "Empty body" - log.error( - "Error-Client: Request URL: %s, Request Body: %s, Response Body: %s", - req.url, - req_body, - response_body, - ) + log.error() raise error return True From b1d08f8f5802abea7beb6cdedc89cab514e79491 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 16:01:27 +0800 Subject: [PATCH 46/71] fix: update imports for merged architecture - Replace get_embedding(llm_settings) with Embeddings().get_embedding() - Remove stray GraphRAGQuery import - Use clean HEAD version of graph_query_node.py --- .../nodes/hugegraph_node/graph_query_node.py | 471 ++++++++++++++++-- .../index_node/build_gremlin_example_index.py | 5 +- 2 files changed, 442 insertions(+), 34 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py index 233e73e8c..7bc9dab69 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py @@ -13,13 +13,63 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, Any +import json from PyCGraph import CStatus +from typing import Dict, Any, Tuple, List, Set, Optional from hugegraph_llm.nodes.base_node import BaseNode -from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.operators.operator_list import OperatorList from hugegraph_llm.utils.log import log +from pyhugegraph.client import PyHugeClient + +# TODO: remove 'as('subj)' step +VERTEX_QUERY_TPL = "g.V({keywords}).limit(8).as('subj').toList()" + +# TODO: we could use a simpler query (like kneighbor-api to get the edges) +# TODO: test with profile()/explain() to speed up the query +VID_QUERY_NEIGHBOR_TPL = """\ +g.V({keywords}) +.repeat( + bothE({edge_labels}).limit({edge_limit}).otherV().dedup() +).times({max_deep}).emit() +.simplePath() +.path() +.by(project('label', 'id', 'props') + .by(label()) + .by(id()) + .by(valueMap().by(unfold())) +) +.by(project('label', 'inV', 'outV', 'props') + .by(label()) + .by(inV().id()) + .by(outV().id()) + .by(valueMap().by(unfold())) +) +.limit({max_items}) +.toList() +""" + +PROPERTY_QUERY_NEIGHBOR_TPL = """\ +g.V().has('{prop}', within({keywords})) +.repeat( + bothE({edge_labels}).limit({edge_limit}).otherV().dedup() +).times({max_deep}).emit() +.simplePath() +.path() +.by(project('label', 'props') + .by(label()) + .by(valueMap().by(unfold())) +) +.by(project('label', 'inV', 'outV', 'props') + .by(label()) + .by(inV().values('{prop}')) + .by(outV().values('{prop}')) + .by(valueMap().by(unfold())) +) +.limit({max_items}) +.toList() +""" class GraphQueryNode(BaseNode): @@ -27,42 +77,387 @@ class GraphQueryNode(BaseNode): Graph query node, responsible for retrieving relevant information from the graph database. """ - graph_rag_query: GraphRAGQuery - def node_init(self): """ Initialize the graph query operator. """ try: - graph_name = huge_settings.graph_name - if not graph_name: - return CStatus(-1, "graph_name is required in wk_input") - - max_deep = self.wk_input.max_deep or 2 - max_graph_items = self.wk_input.max_graph_items or huge_settings.max_graph_items - max_v_prop_len = self.wk_input.max_v_prop_len or 2048 - max_e_prop_len = self.wk_input.max_e_prop_len or 256 - prop_to_match = self.wk_input.prop_to_match - num_gremlin_generate_example = self.wk_input.gremlin_tmpl_num or -1 - gremlin_prompt = self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt - - # Initialize GraphRAGQuery operator - self.graph_rag_query = GraphRAGQuery( - max_deep=max_deep, - max_graph_items=max_graph_items, - max_v_prop_len=max_v_prop_len, - max_e_prop_len=max_e_prop_len, - prop_to_match=prop_to_match, - num_gremlin_generate_example=num_gremlin_generate_example, - gremlin_prompt=gremlin_prompt, + self._client: PyHugeClient = 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, + ) + self._max_deep = self.wk_input.max_deep or 2 + self._max_items = ( + self.wk_input.max_graph_items or huge_settings.max_graph_items + ) + self._prop_to_match = self.wk_input.prop_to_match + self._num_gremlin_generate_example = self.wk_input.gremlin_tmpl_num or -1 + self.gremlin_prompt = ( + self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt ) + self._limit_property = huge_settings.limit_property.lower() == "true" + self._max_v_prop_len = self.wk_input.max_v_prop_len or 2048 + self._max_e_prop_len = self.wk_input.max_e_prop_len or 256 + self._schema = "" + self.operator_list = OperatorList(None, None) return super().node_init() except Exception as e: - log.error("Failed to initialize GraphQueryNode: %s", e) + log.error(f"Failed to initialize GraphQueryNode: {e}") return CStatus(-1, f"GraphQueryNode initialization failed: {e}") + # TODO: move this method to a util file for reuse (remove self param) + def init_client(self, context): + """Initialize the HugeGraph client from context or default settings.""" + # pylint: disable=R0915 (too-many-statements) + if self._client is None: + if isinstance(context.get("graph_client"), PyHugeClient): + self._client = context["graph_client"] + else: + url = context.get("url") or "http://localhost:8080" + graph = context.get("graph") or "hugegraph" + user = context.get("user") or "admin" + pwd = context.get("pwd") or "admin" + gs = context.get("graphspace") or None + self._client = PyHugeClient(url, graph, user, pwd, gs) + assert self._client is not None, "No valid graph to search." + + def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: + query = context["query"] + vertices = context.get("match_vids") + query_embedding = context.get("query_embedding") + + self.operator_list.clear() + self.operator_list.example_index_query( + num_examples=self._num_gremlin_generate_example + ) + gremlin_response = self.operator_list.gremlin_generate_synthesize( + context["simple_schema"], + vertices=vertices, + gremlin_prompt=self.gremlin_prompt, + ).run(query=query, query_embedding=query_embedding) + if self._num_gremlin_generate_example > 0: + gremlin = gremlin_response["result"] + else: + gremlin = gremlin_response["raw_result"] + log.info("Generated gremlin: %s", gremlin) + context["gremlin"] = gremlin + try: + result = self._client.gremlin().exec(gremlin=gremlin)["data"] + if result == [None]: + result = [] + context["graph_result"] = [ + json.dumps(item, ensure_ascii=False) for item in result + ] + if context["graph_result"]: + context["graph_result_flag"] = 1 + context["graph_context_head"] = ( + f"The following are graph query result " + f"from gremlin query `{gremlin}`.\n" + ) + except Exception as e: # pylint: disable=broad-except + log.error(e) + context["graph_result"] = [] + return context + + def _limit_property_query( + self, value: Optional[str], item_type: str + ) -> Optional[str]: + # NOTE: we skip the filter for list/set type (e.g., list of string, add it if needed) + if not self._limit_property or not isinstance(value, str): + return value + + max_len = self._max_v_prop_len if item_type == "v" else self._max_e_prop_len + return value[:max_len] if value else value + + def _process_vertex( + self, + item: Any, + flat_rel: str, + node_cache: Set[str], + prior_edge_str_len: int, + depth: int, + nodes_with_degree: List[str], + use_id_to_match: bool, + v_cache: Set[str], + ) -> Tuple[str, int, int]: + matched_str = ( + item["id"] if use_id_to_match else item["props"][self._prop_to_match] + ) + if matched_str in node_cache: + flat_rel = flat_rel[:-prior_edge_str_len] + return flat_rel, prior_edge_str_len, depth + + node_cache.add(matched_str) + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'v')}" + for k, v in item["props"].items() + if v + ) + + # TODO: we may remove label id or replace with label name + if matched_str in v_cache: + node_str = matched_str + else: + v_cache.add(matched_str) + node_str = f"{item['id']}{{{props_str}}}" + + flat_rel += node_str + nodes_with_degree.append(node_str) + depth += 1 + return flat_rel, prior_edge_str_len, depth + + def _process_edge( + self, + item: Any, + path_str: str, + raw_flat_rel: List[Any], + i: int, + use_id_to_match: bool, + e_cache: Set[Tuple[str, str, str]], + ) -> Tuple[str, int]: + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'e')}" + for k, v in item["props"].items() + if v + ) + props_str = f"{{{props_str}}}" if props_str else "" + prev_matched_str = ( + raw_flat_rel[i - 1]["id"] + if use_id_to_match + else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] + ) + + edge_key = (item["inV"], item["label"], item["outV"]) + if edge_key not in e_cache: + e_cache.add(edge_key) + edge_label = f"{item['label']}{props_str}" + else: + edge_label = item["label"] + + edge_str = ( + f"--[{edge_label}]-->" + if item["outV"] == prev_matched_str + else f"<--[{edge_label}]--" + ) + path_str += edge_str + prior_edge_str_len = len(edge_str) + return path_str, prior_edge_str_len + + def _process_path( + self, + path: Any, + use_id_to_match: bool, + v_cache: Set[str], + e_cache: Set[Tuple[str, str, str]], + ) -> Tuple[str, List[str]]: + flat_rel = "" + raw_flat_rel = path["objects"] + assert len(raw_flat_rel) % 2 == 1, "The length of raw_flat_rel should be odd." + + node_cache = set() + prior_edge_str_len = 0 + depth = 0 + nodes_with_degree = [] + + for i, item in enumerate(raw_flat_rel): + if i % 2 == 0: + # Process each vertex + flat_rel, prior_edge_str_len, depth = self._process_vertex( + item, + flat_rel, + node_cache, + prior_edge_str_len, + depth, + nodes_with_degree, + use_id_to_match, + v_cache, + ) + else: + # Process each edge + flat_rel, prior_edge_str_len = self._process_edge( + item, flat_rel, raw_flat_rel, i, use_id_to_match, e_cache + ) + + return flat_rel, nodes_with_degree + + def _update_vertex_degree_list( + self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str] + ) -> None: + for depth, node_str in enumerate(nodes_with_degree): + if depth >= len(vertex_degree_list): + vertex_degree_list.append(set()) + vertex_degree_list[depth].add(node_str) + + def _format_graph_query_result( + self, query_paths + ) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: + use_id_to_match = self._prop_to_match is None + subgraph = set() + subgraph_with_degree = {} + vertex_degree_list: List[Set[str]] = [] + v_cache: Set[str] = set() + e_cache: Set[Tuple[str, str, str]] = set() + + for path in query_paths: + # 1. Process each path + path_str, vertex_with_degree = self._process_path( + path, use_id_to_match, v_cache, e_cache + ) + subgraph.add(path_str) + subgraph_with_degree[path_str] = vertex_with_degree + # 2. Update vertex degree list + self._update_vertex_degree_list(vertex_degree_list, vertex_with_degree) + + return subgraph, vertex_degree_list, subgraph_with_degree + + def _get_graph_schema(self, refresh: bool = False) -> str: + if self._schema and not refresh: + return self._schema + + schema = self._client.schema() + vertex_schema = schema.getVertexLabels() + edge_schema = schema.getEdgeLabels() + relationships = schema.getRelations() + + self._schema = ( + f"Vertex properties: {vertex_schema}\n" + f"Edge properties: {edge_schema}\n" + f"Relationships: {relationships}\n" + ) + log.debug("Link(Relation): %s", relationships) + return self._schema + + @staticmethod + def _extract_label_names( + source: str, head: str = "name: ", tail: str = ", " + ) -> List[str]: + result = [] + for s in source.split(head): + end = s.find(tail) + label = s[:end] + if label: + result.append(label) + return result + + def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: + schema = self._get_graph_schema() + vertex_props_str, edge_props_str = schema.split("\n")[:2] + # TODO: rename to vertex (also need update in the schema) + vertex_props_str = ( + vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") + ) + edge_props_str = ( + edge_props_str[len("Edge properties: ") :].strip("[").strip("]") + ) + vertex_labels = self._extract_label_names(vertex_props_str) + edge_labels = self._extract_label_names(edge_props_str) + return vertex_labels, edge_labels + + def _format_graph_from_vertex(self, query_result: List[Any]) -> Set[str]: + knowledge = set() + for item in query_result: + props_str = ", ".join(f"{k}: {v}" for k, v in item["properties"].items()) + node_str = f"{item['id']}{{{props_str}}}" + knowledge.add(node_str) + return knowledge + + def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: + # 1. Extract params from context + matched_vids = context.get("match_vids") + if isinstance(context.get("max_deep"), int): + self._max_deep = context["max_deep"] + if isinstance(context.get("max_items"), int): + self._max_items = context["max_items"] + if isinstance(context.get("prop_to_match"), str): + self._prop_to_match = context["prop_to_match"] + + # 2. Extract edge_labels from graph schema + _, edge_labels = self._extract_labels_from_schema() + edge_labels_str = ",".join("'" + label + "'" for label in edge_labels) + # TODO: enhance the limit logic later + edge_limit_amount = len(edge_labels) * huge_settings.edge_limit_pre_label + + use_id_to_match = self._prop_to_match is None + if use_id_to_match: + if not matched_vids: + return context + + gremlin_query = VERTEX_QUERY_TPL.format(keywords=matched_vids) + vertexes = self._client.gremlin().exec(gremlin=gremlin_query)["data"] + log.debug("Vids gremlin query: %s", gremlin_query) + + vertex_knowledge = self._format_graph_from_vertex(query_result=vertexes) + paths: List[Any] = [] + # TODO: use generator or asyncio to speed up the query logic + for matched_vid in matched_vids: + gremlin_query = VID_QUERY_NEIGHBOR_TPL.format( + keywords=f"'{matched_vid}'", + max_deep=self._max_deep, + edge_labels=edge_labels_str, + edge_limit=edge_limit_amount, + max_items=self._max_items, + ) + log.debug("Kneighbor gremlin query: %s", gremlin_query) + paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) + + ( + graph_chain_knowledge, + vertex_degree_list, + knowledge_with_degree, + ) = self._format_graph_query_result(query_paths=paths) + + # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) + if not graph_chain_knowledge: + graph_chain_knowledge.update(vertex_knowledge) + if vertex_degree_list: + vertex_degree_list[0].update(vertex_knowledge) + else: + vertex_degree_list.append(vertex_knowledge) + else: + # WARN: When will the query enter here? + keywords = context.get("keywords") + assert keywords, "No related property(keywords) for graph query." + keywords_str = ",".join("'" + kw + "'" for kw in keywords) + gremlin_query = PROPERTY_QUERY_NEIGHBOR_TPL.format( + prop=self._prop_to_match, + keywords=keywords_str, + edge_labels=edge_labels_str, + edge_limit=edge_limit_amount, + max_deep=self._max_deep, + max_items=self._max_items, + ) + log.warning( + "Unable to find vid, downgraded to property query, please confirm if it meets expectation." + ) + + paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)[ + "data" + ] + ( + graph_chain_knowledge, + vertex_degree_list, + knowledge_with_degree, + ) = self._format_graph_query_result(query_paths=paths) + + context["graph_result"] = list(graph_chain_knowledge) + if context["graph_result"]: + context["graph_result_flag"] = 0 + context["vertex_degree_list"] = [ + list(vertex_degree) for vertex_degree in vertex_degree_list + ] + context["knowledge_with_degree"] = knowledge_with_degree + context["graph_context_head"] = ( + f"The following are graph knowledge in {self._max_deep} depth, e.g:\n" + "`vertexA--[links]-->vertexB<--[links]--vertexC ...`" + "extracted based on key entities as subject:\n" + ) + return context + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: """ Execute the graph query operation. @@ -76,16 +471,30 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: return data_json # Execute the graph query (assuming schema and semantic query have been completed in previous nodes) - graph_result = self.graph_rag_query.run(data_json) - data_json.update(graph_result) + self.init_client(data_json) + + # initial flag: -1 means no result, 0 means subgraph query, 1 means gremlin query + data_json["graph_result_flag"] = -1 + # 1. Try to perform a query based on the generated gremlin + if self._num_gremlin_generate_example >= 0: + data_json = self._gremlin_generate_query(data_json) + # 2. Try to perform a query based on subgraph-search if the previous query failed + if not data_json.get("graph_result"): + data_json = self._subgraph_query(data_json) + + if data_json.get("graph_result"): + log.debug( + "Knowledge from Graph:\n%s", "\n".join(data_json["graph_result"]) + ) + else: + log.debug("No Knowledge Extracted from Graph") log.info( - "Graph query completed, found %d results", - len(data_json.get("graph_result", [])), + f"Graph query completed, found {len(data_json.get('graph_result', []))} results" ) return data_json except Exception as e: - log.error("Graph query failed: %s", e) + log.error(f"Graph query failed: {e}") return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py index 8772959d7..5a729a94b 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py @@ -15,8 +15,7 @@ from PyCGraph import CStatus -from hugegraph_llm.config import llm_settings -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_gremlin_example_index import ( BuildGremlinExampleIndex, @@ -35,7 +34,7 @@ def node_init(self): examples = self.wk_input.examples self.build_gremlin_example_index_op = BuildGremlinExampleIndex( - get_embedding(llm_settings), examples + Embeddings().get_embedding(), examples ) return super().node_init() From 69283381d10ec9d917c1d64e46998686e776bce5 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 16:13:38 +0800 Subject: [PATCH 47/71] fix: update graph_index_utils to use modular vector stores - Replace VectorIndex with FaissVectorIndex - Remove deprecated kg_construction_task import from vector_index_utils - Align with new vector store architecture --- hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py | 6 +++--- hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py | 1 - 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index a0e376ac6..5f50bcc50 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -29,7 +29,7 @@ from .log import log from .vector_index_utils import read_documents from ..config import resource_path, huge_settings, llm_settings -from ..indices.vector_index import VectorIndex +from ..indices.vector_index.faiss_vector_store import FaissVectorIndex from ..models.embeddings.init_embedding import Embeddings @@ -50,10 +50,10 @@ def clean_all_graph_index(): llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None), ) - VectorIndex.clean( + FaissVectorIndex.clean( str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix ) - VectorIndex.clean( + FaissVectorIndex.clean( str(os.path.join(resource_path, folder_name, "gremlin_examples")), filename_prefix, ) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 9468f823b..b01a43a2d 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -26,7 +26,6 @@ from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.kg_construction_task import KgBuilder from hugegraph_llm.utils.hugegraph_utils import get_hg_client From bd89ab637af55e364e455f709f36d6d3b8989b5a Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 16:15:13 +0800 Subject: [PATCH 48/71] fix: use HEAD version of demo UI blocks - Replace rag_block.py and text2gremlin_block.py with HEAD versions - Remove deprecated RAGPipeline imports - Use SchedulerSingleton for all flow execution - App module now imports successfully --- .../hugegraph_llm/demo/rag_demo/rag_block.py | 196 ++++++++-------- .../demo/rag_demo/text2gremlin_block.py | 210 +++++++----------- 2 files changed, 183 insertions(+), 223 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 5482cdd87..60ca6ae55 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -18,20 +18,14 @@ # pylint: disable=E1101 import os -from typing import Any, AsyncGenerator, Literal, Optional, Tuple +from typing import AsyncGenerator, Tuple, Literal, Optional import gradio as gr +from hugegraph_llm.flows.scheduler import SchedulerSingleton import pandas as pd +from gradio.utils import NamedString -from hugegraph_llm.config import ( - huge_settings, - index_settings, - llm_settings, - prompt, - resource_path, -) -from hugegraph_llm.operators.graph_rag_task import RAGPipeline -from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize +from hugegraph_llm.config import resource_path, prompt, llm_settings from hugegraph_llm.utils.decorators import with_task_id from hugegraph_llm.utils.log import log @@ -73,49 +67,51 @@ def rag_answer( gr.Warning("Please select at least one generate mode.") return "", "", "", "" - rag = RAGPipeline() - if vector_search: - rag.query_vector_index(vector_index_str=index_settings.cur_vector_index) - if graph_search: - rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid( - vector_index_str=index_settings.cur_vector_index, - vector_dis_threshold=vector_dis_threshold, - topk_per_keyword=topk_per_keyword, - ).import_schema(huge_settings.graph_name).query_graphdb( - num_gremlin_generate_example=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - max_graph_items=max_graph_items, - ) - # TODO: add more user-defined search strategies - rag.merge_dedup_rerank( - graph_ratio=graph_ratio, - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - topk_return_results=topk_return_results, - ) - rag.synthesize_answer( - raw_answer, - vector_only_answer, - graph_only_answer, - graph_vector_answer, - answer_prompt, - ) - + scheduler = SchedulerSingleton.get_instance() try: - context = rag.run( - verbose=True, + # Select workflow by mode to avoid fetching the wrong pipeline from the pool + if graph_vector_answer or (graph_only_answer and vector_only_answer): + flow_key = "rag_graph_vector" + elif vector_only_answer: + flow_key = "rag_vector_only" + elif graph_only_answer: + flow_key = "rag_graph_only" + elif raw_answer: + flow_key = "rag_raw" + else: + raise RuntimeError("Unsupported flow type") + + res = scheduler.schedule_flow( + flow_key, query=text, vector_search=vector_search, graph_search=graph_search, + raw_answer=raw_answer, + vector_only_answer=vector_only_answer, + graph_only_answer=graph_only_answer, + graph_vector_answer=graph_vector_answer, + graph_ratio=graph_ratio, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + answer_prompt=answer_prompt, + keywords_extract_prompt=keywords_extract_prompt, + gremlin_tmpl_num=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, max_graph_items=max_graph_items, + topk_return_results=topk_return_results, + vector_dis_threshold=vector_dis_threshold, + topk_per_keyword=topk_per_keyword, ) - if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") + if res.get("switch_to_bleu"): + gr.Warning( + "Online reranker fails, automatically switches to local bleu rerank." + ) return ( - context.get("raw_answer", ""), - context.get("vector_only_answer", ""), - context.get("graph_only_answer", ""), - context.get("graph_vector_answer", ""), + res.get("raw_answer", ""), + res.get("vector_only_answer", ""), + res.get("graph_only_answer", ""), + res.get("graph_vector_answer", ""), ) except ValueError as e: log.critical(e) @@ -189,47 +185,47 @@ async def rag_answer_streaming( yield "", "", "", "" return - rag = RAGPipeline() - if vector_search: - rag.query_vector_index(vector_index_str=index_settings.cur_vector_index) - if graph_search: - rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid( - vector_index_str=index_settings.cur_vector_index - ).import_schema(huge_settings.graph_name).query_graphdb( - num_gremlin_generate_example=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - ) - rag.merge_dedup_rerank( - graph_ratio, - rerank_method, - near_neighbor_first, - ) - # rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) - try: - context = rag.run( - verbose=True, + # Select the specific streaming workflow + scheduler = SchedulerSingleton.get_instance() + if graph_vector_answer or (graph_only_answer and vector_only_answer): + flow_key = "rag_graph_vector" + elif vector_only_answer: + flow_key = "rag_vector_only" + elif graph_only_answer: + flow_key = "rag_graph_only" + elif raw_answer: + flow_key = "rag_raw" + else: + raise RuntimeError("Unsupported flow type") + + async for res in scheduler.schedule_stream_flow( + flow_key, query=text, vector_search=vector_search, graph_search=graph_search, - ) - if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") - answer_synthesize = AnswerSynthesize( raw_answer=raw_answer, vector_only_answer=vector_only_answer, graph_only_answer=graph_only_answer, graph_vector_answer=graph_vector_answer, - prompt_template=answer_prompt, - ) - async for context in answer_synthesize.run_streaming(context): - if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") + graph_ratio=graph_ratio, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + answer_prompt=answer_prompt, + keywords_extract_prompt=keywords_extract_prompt, + gremlin_tmpl_num=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, + ): + if res.get("switch_to_bleu"): + gr.Warning( + "Online reranker fails, automatically switches to local bleu rerank." + ) yield ( - context.get("raw_answer", ""), - context.get("vector_only_answer", ""), - context.get("graph_only_answer", ""), - context.get("graph_vector_answer", ""), + res.get("raw_answer", ""), + res.get("vector_only_answer", ""), + res.get("graph_only_answer", ""), + res.get("graph_vector_answer", ""), ) except ValueError as e: log.critical(e) @@ -294,7 +290,9 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): - raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") + raw_radio = gr.Radio( + choices=[True, False], value=False, label="Basic LLM Answer" + ) vector_only_radio = gr.Radio( choices=[True, False], value=False, label="Vector-only Answer" ) @@ -365,7 +363,7 @@ def toggle_slider(enable): """## 2. (Batch) Back-testing ) > 1. Download the template file & fill in the questions you want to test. > 2. Upload the file & click the button to generate answers. (Preview shows the first 40 lines) - > 3. The answer options are the same as the above RAG/Q&A frame + > 3. The answer options are the same as the above RAG/Q&A frame """ ) tests_df_headers = [ @@ -379,9 +377,11 @@ def toggle_slider(enable): # FIXME: "demo" might conflict with the graph name, it should be modified. answers_path = os.path.join(resource_path, "demo", "questions_answers.xlsx") questions_path = os.path.join(resource_path, "demo", "questions.xlsx") - questions_template_path = os.path.join(resource_path, "demo", "questions_template.xlsx") + questions_template_path = os.path.join( + resource_path, "demo", "questions_template.xlsx" + ) - def read_file_to_excel(file: Any, line_count: Optional[int] = None): + def read_file_to_excel(file: NamedString, line_count: Optional[int] = None): df = None if not file: return pd.DataFrame(), 1 @@ -389,15 +389,15 @@ def read_file_to_excel(file: Any, line_count: Optional[int] = None): df = pd.read_excel(file.name, nrows=line_count) if file else pd.DataFrame() elif file.name.endswith(".csv"): df = pd.read_csv(file.name, nrows=line_count) if file else pd.DataFrame() - df.to_excel(questions_path, index=False) # type:ignore - if df.empty: # type:ignore + df.to_excel(questions_path, index=False) + if df.empty: df = pd.DataFrame([[""] * len(tests_df_headers)], columns=tests_df_headers) else: - df.columns = tests_df_headers # type:ignore + df.columns = tests_df_headers # truncate the dataframe if it's too long - if len(df) > 40: # type:ignore - return df.head(40), 40 # type:ignore - return df, len(df) # type:ignore + if len(df) > 40: + return df.head(40), 40 + return df, len(df) def change_showing_excel(line_count): if os.path.exists(answers_path): @@ -459,12 +459,18 @@ def several_rag_answer( file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" ) with gr.Column(): - test_template_file = os.path.join(resource_path, "demo", "questions_template.xlsx") + test_template_file = os.path.join( + resource_path, "demo", "questions_template.xlsx" + ) gr.File(value=test_template_file, label="Download Template File") - answer_max_line_count = gr.Number(1, label="Max Lines To Show", minimum=1, maximum=40) + answer_max_line_count = gr.Number( + 1, label="Max Lines To Show", minimum=1, maximum=40 + ) answers_btn = gr.Button("Generate Answer (Batch)", variant="primary") # TODO: Set individual progress bars for dataframe - qa_dataframe = gr.DataFrame(label="Questions & Answers (Preview)", headers=tests_df_headers) + qa_dataframe = gr.DataFrame( + label="Questions & Answers (Preview)", headers=tests_df_headers + ) answers_btn.click( several_rag_answer, inputs=[ @@ -482,8 +488,12 @@ def several_rag_answer( ], outputs=[qa_dataframe, gr.File(label="Download Answered File", min_width=40)], ) - questions_file.change(read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count]) - answer_max_line_count.change(change_showing_excel, answer_max_line_count, qa_dataframe) + questions_file.change( + read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count] + ) + answer_max_line_count.change( + change_showing_excel, answer_max_line_count, qa_dataframe + ) return ( inp, answer_prompt_input, diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 8fbb01c25..04aef1c77 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -17,23 +17,18 @@ import json import os -from dataclasses import dataclass from datetime import datetime -from typing import Any, Dict, List, Literal, Optional, Tuple +from dataclasses import dataclass +from typing import Any, Tuple, Dict, Literal, Optional, List import gradio as gr import pandas as pd -from hugegraph_llm.config import huge_settings, index_settings, prompt, resource_path -from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.graph_rag_task import RAGPipeline -from hugegraph_llm.operators.gremlin_generate_task import GremlinGenerator -from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager +from hugegraph_llm.config import prompt, resource_path, huge_settings from hugegraph_llm.utils.embedding_utils import get_index_folder_name from hugegraph_llm.utils.hugegraph_utils import run_gremlin_query from hugegraph_llm.utils.log import log -from hugegraph_llm.utils.vector_index_utils import get_vector_index_class +from hugegraph_llm.flows.scheduler import SchedulerSingleton @dataclass @@ -86,8 +81,12 @@ def store_schema(schema, question, gremlin_prompt): def build_example_vector_index(temp_file) -> dict: - vector_index = get_vector_index_class(index_settings.cur_vector_index) - assert vector_index, "vector db name is error" + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + index_path = os.path.join(resource_path, folder_name, "gremlin_examples") + if not os.path.exists(index_path): + os.makedirs(index_path) if temp_file is None: full_path = os.path.join(resource_path, "demo", "text2gremlin.csv") else: @@ -96,12 +95,12 @@ def build_example_vector_index(temp_file) -> dict: timestamp = datetime.now().strftime("%Y%m%d%H%M%S") _, file_name = os.path.split(f"{name}_{timestamp}{ext}") log.info("Copying file to: %s", file_name) - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - target_file = os.path.join(resource_path, folder_name, "gremlin_examples", file_name) + target_file = os.path.join( + resource_path, folder_name, "gremlin_examples", file_name + ) try: import shutil - os.makedirs(os.path.dirname(target_file), exist_ok=True) shutil.copy2(full_path, target_file) log.info("Successfully copied file to: %s", target_file) except (OSError, IOError) as e: @@ -116,11 +115,10 @@ def build_example_vector_index(temp_file) -> dict: else: log.critical("Unsupported file format. Please input a JSON or CSV file.") return {"error": "Unsupported file format. Please input a JSON or CSV file."} - builder = GremlinGenerator( - llm=LLMs().get_text2gql_llm(), - embedding=Embeddings().get_embedding(), + + return SchedulerSingleton.get_instance().schedule_flow( + "build_examples_index", examples ) - return builder.example_index_build(examples, vector_index=vector_index).run() def _process_schema(schema, generator, sm): @@ -182,56 +180,18 @@ def _execute_queries(context, output_types): context["raw_exec_res"] = "" -def gremlin_generate( - inp, - example_num, - schema, - gremlin_prompt, - requested_outputs: Optional[List[str]] = None, -) -> GremlinResult: - vector_index = get_vector_index_class(index_settings.cur_vector_index) - generator = GremlinGenerator( - llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() - ) - sm = SchemaManager(graph_name=schema) - - processed_schema, short_schema = _process_schema(schema, generator, sm) - if processed_schema is None and short_schema is None: - return GremlinResult.error("Invalid JSON schema, please check the format carefully.") - - updated_schema = sm.simple_schema(processed_schema) if short_schema else processed_schema - store_schema(str(updated_schema), inp, gremlin_prompt) - - output_types = _configure_output_types(requested_outputs) - - context = ( - generator.example_index_query(example_num, vector_index) - .gremlin_generate_synthesize(updated_schema, gremlin_prompt) - .run(query=inp) - ) - - _execute_queries(context, output_types) - - match_result = json.dumps( - context.get("match_result", "No Results"), ensure_ascii=False, indent=2 - ) - return GremlinResult.success_result( - match_result=match_result, - template_gremlin=context["result"], - raw_gremlin=context["raw_result"], - template_exec=context["template_exec_res"], - raw_exec=context["raw_exec_res"], - ) - - def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: - mini_schema = {} # type: ignore + mini_schema = {} # Add necessary vertexlabels items (3) if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} + new_vertex = { + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex + } mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) @@ -250,17 +210,40 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: def gremlin_generate_for_ui(inp, example_num, schema, gremlin_prompt): """UI wrapper for gremlin_generate that returns tuple for Gradio compatibility""" - result = gremlin_generate(inp, example_num, schema, gremlin_prompt) - - if not result.success: - return result.match_result, "", "", "", "" + # Execute via scheduler + try: + res = SchedulerSingleton.get_instance().schedule_flow( + "text2gremlin", + inp, + int(example_num) if isinstance(example_num, (int, float, str)) else 2, + schema, + gremlin_prompt, + [ + "match_result", + "template_gremlin", + "raw_gremlin", + "template_execution_result", + "raw_execution_result", + ], + ) + except Exception as e: # pylint: disable=broad-except + log.error("UI text2gremlin error: %s", e) + return json.dumps({"error": str(e)}, ensure_ascii=False), "", "", "", "" + + # Backward-compatible mapping for outputs + match_result = res.get("match_result", []) + match_result_str = ( + json.dumps(match_result, ensure_ascii=False, indent=2) + if isinstance(match_result, (list, dict)) + else str(match_result) + ) return ( - result.match_result, - result.template_gremlin or "", - result.raw_gremlin or "", - result.template_exec_result or "", - result.raw_exec_result or "", + match_result_str, + res.get("template_gremlin", "") or "", + res.get("raw_gremlin", "") or "", + res.get("template_execution_result", "") or "", + res.get("raw_execution_result", "") or "", ) @@ -281,7 +264,6 @@ def create_text2gremlin_block() -> Tuple: out = gr.Textbox(label="Result Message") with gr.Row(): btn = gr.Button("Build Example Vector Index", variant="primary") - btn.click(build_example_vector_index, inputs=[file], outputs=[out]) # pylint: disable=no-member gr.Markdown("## Nature Language To Gremlin") @@ -297,8 +279,12 @@ def create_text2gremlin_block() -> Tuple: language="javascript", elem_classes="code-container-show", ) - initialized_out = gr.Textbox(label="Gremlin With Template", show_copy_button=True) - raw_out = gr.Textbox(label="Gremlin Without Template", show_copy_button=True) + initialized_out = gr.Textbox( + label="Gremlin With Template", show_copy_button=True + ) + raw_out = gr.Textbox( + label="Gremlin Without Template", show_copy_button=True + ) tmpl_exec_out = gr.Code( label="Query With Template Output", language="json", @@ -350,25 +336,21 @@ def graph_rag_recall( get_vertex_only: bool = False, ) -> dict: store_schema(prompt.text2gql_graph_schema, query, gremlin_prompt) - rag = RAGPipeline() - rag.extract_keywords().keywords_to_vid( - vector_index_str=index_settings.cur_vector_index, + context = SchedulerSingleton.get_instance().schedule_flow( + "rag_graph_only", + query=query, + gremlin_tmpl_num=gremlin_tmpl_num, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + gremlin_prompt=gremlin_prompt, + max_graph_items=max_graph_items, + topk_return_results=topk_return_results, vector_dis_threshold=vector_dis_threshold, topk_per_keyword=topk_per_keyword, + is_graph_rag_recall=True, + is_vector_only=get_vertex_only, ) - - if not get_vertex_only: - rag.import_schema(huge_settings.graph_name).query_graphdb( - num_gremlin_generate_example=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - max_graph_items=max_graph_items, - ).merge_dedup_rerank( - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - custom_related_information=custom_related_information, - topk_return_results=topk_return_results, - ) - context = rag.run(verbose=True, query=query, graph_search=True) return context @@ -379,45 +361,13 @@ def gremlin_generate_selective( gremlin_prompt_input: str, requested_outputs: Optional[List[str]] = None, ) -> Dict[str, Any]: - """ - Wraps the gremlin_generate function to return a dictionary of outputs - based on the requested_outputs list of strings. - """ - output_keys = [ - "match_result", - "template_gremlin", - "raw_gremlin", - "template_execution_result", - "raw_execution_result", - ] - if not requested_outputs: # None or empty list - requested_outputs = output_keys - - result = gremlin_generate( - inp, example_num, schema_input, gremlin_prompt_input, requested_outputs + response_dict = SchedulerSingleton.get_instance().schedule_flow( + "text2gremlin", + inp, + example_num, + schema_input, + gremlin_prompt_input, + requested_outputs, ) - outputs_dict: Dict[str, Any] = {} - - if not result.success: - # Handle error case - if "match_result" in requested_outputs: - outputs_dict["match_result"] = result.match_result - if result.error_message: - outputs_dict["error_detail"] = result.error_message - return outputs_dict - - # Handle successful case - output_mapping = { - "match_result": result.match_result, - "template_gremlin": result.template_gremlin, - "raw_gremlin": result.raw_gremlin, - "template_execution_result": result.template_exec_result, - "raw_execution_result": result.raw_exec_result, - } - - for key in requested_outputs: - if key in output_mapping: - outputs_dict[key] = output_mapping[key] - - return outputs_dict + return response_dict From 3a7591676bc986833fcf7116a45225968b0a913f Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 16:17:36 +0800 Subject: [PATCH 49/71] fix: add missing vector_index parameter to BuildSemanticIndex - Update BuildSemanticIndexNode to use get_vector_index_class() - Update operator_list.py build_vertex_id_semantic_index() signature - Pass vector_index based on index_settings.cur_vector_index config - Fixes: BuildSemanticIndex.__init__() missing 1 required positional argument --- .../hugegraph_llm/nodes/index_node/build_semantic_index.py | 5 ++++- hugegraph-llm/src/hugegraph_llm/operators/operator_list.py | 4 ++-- 2 files changed, 6 insertions(+), 3 deletions(-) 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 6a4424f38..cc02bbe5d 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 @@ -13,10 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from hugegraph_llm.config import index_settings from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.utils.vector_index_utils import get_vector_index_class class BuildSemanticIndexNode(BaseNode): @@ -25,7 +27,8 @@ class BuildSemanticIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): - self.build_semantic_index_op = BuildSemanticIndex(Embeddings().get_embedding()) + vector_index = get_vector_index_class(index_settings.cur_vector_index) + self.build_semantic_index_op = BuildSemanticIndex(Embeddings().get_embedding(), vector_index) return super().node_init() def operator_schedule(self, data_json): diff --git a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py index 702258ace..ec7fb9522 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py @@ -134,8 +134,8 @@ def commit_to_hugegraph(self): self.operators.append(Commit2Graph()) return self - def build_vertex_id_semantic_index(self): - self.operators.append(BuildSemanticIndex(self.embedding)) + def build_vertex_id_semantic_index(self, vector_index): + self.operators.append(BuildSemanticIndex(self.embedding, vector_index)) return self def build_vector_index(self): From a10730c5402ac378f096ba57f6366b5b28189c7b Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 16:19:59 +0800 Subject: [PATCH 50/71] fix: remove KgBuilder and resolve circular import - Replace KgBuilder usage with SchedulerSingleton in build_vector_index() - Add SchedulerSingleton import to vector_index_utils.py - Use lazy import in BuildSemanticIndexNode to avoid circular dependency - Remove unused LLMs and get_hg_client imports - Fixes: NameError: name 'KgBuilder' is not defined --- .../nodes/index_node/build_semantic_index.py | 4 +++- .../src/hugegraph_llm/utils/vector_index_utils.py | 9 +++------ 2 files changed, 6 insertions(+), 7 deletions(-) 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 cc02bbe5d..7b6335711 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 @@ -18,7 +18,6 @@ from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState -from hugegraph_llm.utils.vector_index_utils import get_vector_index_class class BuildSemanticIndexNode(BaseNode): @@ -27,6 +26,9 @@ class BuildSemanticIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): + # Lazy import to avoid circular dependency + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel + vector_index = get_vector_index_class(index_settings.cur_vector_index) self.build_semantic_index_op = BuildSemanticIndex(Embeddings().get_embedding(), vector_index) return super().node_init() diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index b01a43a2d..bb9536d25 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -22,11 +22,10 @@ import gradio as gr from hugegraph_llm.config import huge_settings, index_settings +from hugegraph_llm.flows.scheduler import SchedulerSingleton from hugegraph_llm.indices.vector_index.base import VectorStoreBase from hugegraph_llm.indices.vector_index.faiss_vector_store import FaissVectorIndex from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.utils.hugegraph_utils import get_hg_client def read_documents(input_file, input_text): @@ -80,13 +79,11 @@ def clean_vector_index(): def build_vector_index(input_file, input_text): - vector_index = get_vector_index_class(index_settings.cur_vector_index) if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - context = builder.chunk_split(texts, "paragraph", "zh").build_vector_index(vector_index).run() - return json.dumps(context, ensure_ascii=False, indent=2) + scheduler = SchedulerSingleton.get_instance() + return scheduler.schedule_flow("build_vector_index", texts) def get_vector_index_class(vector_index_str: str) -> Type[VectorStoreBase]: From 77343c268e0529ea5712d4f39a25a6df9e14da00 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 16:23:28 +0800 Subject: [PATCH 51/71] fix: add missing vector_index parameter to BuildVectorIndex - Update BuildVectorIndexNode to use get_vector_index_class() - Update operator_list.py build_vector_index() signature - Use lazy import to avoid circular dependency - Fixes: BuildVectorIndex.__init__() missing 1 required positional argument --- .../hugegraph_llm/nodes/index_node/build_vector_index.py | 7 ++++++- hugegraph-llm/src/hugegraph_llm/operators/operator_list.py | 4 ++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py index 28f2cb041..745c5d32b 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from hugegraph_llm.config import index_settings from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex @@ -25,7 +26,11 @@ class BuildVectorIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): - self.build_vector_index_op = BuildVectorIndex(Embeddings().get_embedding()) + # Lazy import to avoid circular dependency + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel + + vector_index = get_vector_index_class(index_settings.cur_vector_index) + self.build_vector_index_op = BuildVectorIndex(Embeddings().get_embedding(), vector_index) return super().node_init() def operator_schedule(self, data_json): diff --git a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py index ec7fb9522..8cf964406 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py @@ -138,8 +138,8 @@ def build_vertex_id_semantic_index(self, vector_index): self.operators.append(BuildSemanticIndex(self.embedding, vector_index)) return self - def build_vector_index(self): - self.operators.append(BuildVectorIndex(self.embedding)) + def build_vector_index(self, vector_index): + self.operators.append(BuildVectorIndex(self.embedding, vector_index)) return self def extract_word(self, text: Optional[str] = None, language: str = "english"): From d212091948ebb714359ee650fa6a242d55ff6e73 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 16:26:27 +0800 Subject: [PATCH 52/71] fix: add missing vector_index parameter to VectorIndexQuery - Update VectorQueryNode to use get_vector_index_class() - Update operator_list.py query_vector_index() signature - Use lazy import to avoid circular dependency - Fixes: VectorIndexQuery.__init__() missing 1 required positional argument --- .../hugegraph_llm/nodes/index_node/vector_query_node.py | 7 ++++++- hugegraph-llm/src/hugegraph_llm/operators/operator_list.py | 4 +++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py index f9af6c49d..ed95fbe9c 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py @@ -14,6 +14,7 @@ # limitations under the License. from typing import Dict, Any +from hugegraph_llm.config import index_settings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery from hugegraph_llm.models.embeddings.init_embedding import Embeddings @@ -32,11 +33,15 @@ def node_init(self): Initialize the vector query operator """ try: + # Lazy import to avoid circular dependency + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel + # 从 wk_input 中读取用户配置参数 + vector_index = get_vector_index_class(index_settings.cur_vector_index) embedding = Embeddings().get_embedding() max_items = self.wk_input.max_items if self.wk_input.max_items is not None else 3 - self.operator = VectorIndexQuery(embedding=embedding, topk=max_items) + self.operator = VectorIndexQuery(vector_index=vector_index, embedding=embedding, topk=max_items) return super().node_init() except Exception as e: log.error("Failed to initialize VectorQueryNode: %s", e) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py index 8cf964406..f8a27a6de 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py @@ -205,15 +205,17 @@ def keywords_to_vid( ) return self - def query_vector_index(self, max_items: int = 3): + def query_vector_index(self, vector_index, max_items: int = 3): """ Add a vector index query operator to the pipeline. + :param vector_index: Vector index class to use. :param max_items: Maximum number of items to retrieve. :return: Self-instance for chaining. """ self.operators.append( VectorIndexQuery( + vector_index=vector_index, embedding=self.embedding, topk=max_items, ) From c5e4bb08472a87e7b8db9a161798599d40e642c0 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 16:28:14 +0800 Subject: [PATCH 53/71] fix: add missing vector_index parameter to SemanticIdQuery - Update SemanticIdQueryNode to use get_vector_index_class() - Update operator_list.py query_semantic_id() signature - Use lazy import to avoid circular dependency - Fixes: SemanticIdQuery.__init__() missing 1 required positional argument --- .../nodes/index_node/semantic_id_query_node.py | 7 ++++++- hugegraph-llm/src/hugegraph_llm/operators/operator_list.py | 3 +++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py index ba5261c59..7156eb0c7 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -19,7 +19,7 @@ from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.semantic_id_query import SemanticIdQuery from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.config import huge_settings +from hugegraph_llm.config import huge_settings, index_settings from hugegraph_llm.utils.log import log @@ -35,10 +35,14 @@ def node_init(self): Initialize the semantic ID query operator. """ try: + # Lazy import to avoid circular dependency + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel + graph_name = huge_settings.graph_name if not graph_name: return CStatus(-1, "graph_name is required in wk_input") + vector_index = get_vector_index_class(index_settings.cur_vector_index) embedding = Embeddings().get_embedding() by = self.wk_input.semantic_by or "keywords" topk_per_keyword = self.wk_input.topk_per_keyword or huge_settings.topk_per_keyword @@ -52,6 +56,7 @@ def node_init(self): # Initialize the semantic ID query operator self.semantic_id_query = SemanticIdQuery( embedding=embedding, + vector_index=vector_index, by=by, topk_per_keyword=topk_per_keyword, topk_per_query=topk_per_query, diff --git a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py index f8a27a6de..acce7d1b5 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py @@ -181,6 +181,7 @@ def extract_keywords( def keywords_to_vid( self, + vector_index, by: Literal["query", "keywords"] = "keywords", topk_per_keyword: int = huge_settings.topk_per_keyword, topk_per_query: int = 10, @@ -188,6 +189,7 @@ def keywords_to_vid( ): """ Add a semantic ID query operator to the pipeline. + :param vector_index: Vector index class to use. :param by: Match by query or keywords. :param topk_per_keyword: Top K results per keyword. :param topk_per_query: Top K results per query. @@ -197,6 +199,7 @@ def keywords_to_vid( self.operators.append( SemanticIdQuery( embedding=self.embedding, + vector_index=vector_index, by=by, topk_per_keyword=topk_per_keyword, topk_per_query=topk_per_query, From ccc6b9a79fd94a0bc3fd5d5da19b036c4aded100 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 16:39:01 +0800 Subject: [PATCH 54/71] fix: add missing vector_index parameter to BuildGremlinExampleIndex - Update BuildGremlinExampleIndexNode to use get_vector_index_class() - Update operator_list.py example_index_build() signature - Use lazy import to avoid circular dependency - Fixes: BuildGremlinExampleIndex.__init__() missing 1 required positional argument --- .../nodes/index_node/build_gremlin_example_index.py | 7 ++++++- hugegraph-llm/src/hugegraph_llm/operators/operator_list.py | 4 ++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py index 5a729a94b..d00a35048 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py @@ -15,6 +15,7 @@ from PyCGraph import CStatus +from hugegraph_llm.config import index_settings from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_gremlin_example_index import ( @@ -29,12 +30,16 @@ class BuildGremlinExampleIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): + # Lazy import to avoid circular dependency + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel + if not self.wk_input.examples: return CStatus(-1, "examples is required in BuildGremlinExampleIndexNode") examples = self.wk_input.examples + vector_index = get_vector_index_class(index_settings.cur_vector_index) self.build_gremlin_example_index_op = BuildGremlinExampleIndex( - Embeddings().get_embedding(), examples + Embeddings().get_embedding(), examples, vector_index ) return super().node_init() diff --git a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py index acce7d1b5..0f815b929 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py @@ -64,8 +64,8 @@ def clear(self): self.operators = [] return self - def example_index_build(self, examples): - self.operators.append(BuildGremlinExampleIndex(self.embedding, examples)) + def example_index_build(self, examples, vector_index): + self.operators.append(BuildGremlinExampleIndex(self.embedding, examples, vector_index)) return self def import_schema( From 31de3b46082f7a15f5169596a13a0e55912007b1 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Fri, 17 Oct 2025 16:40:14 +0800 Subject: [PATCH 55/71] fix: add missing vector_index parameter to GremlinExampleIndexQuery - Update GremlinExampleIndexQueryNode to use get_vector_index_class() - Update operator_list.py example_index_query() signature - Use lazy import to avoid circular dependency - Fixes: GremlinExampleIndexQuery.__init__() missing 1 required positional argument --- .../nodes/index_node/gremlin_example_index_query.py | 7 ++++++- hugegraph-llm/src/hugegraph_llm/operators/operator_list.py | 4 ++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py index 8b9e0db4d..2c801f726 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -19,6 +19,7 @@ from PyCGraph import CStatus +from hugegraph_llm.config import index_settings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.gremlin_example_index_query import ( GremlinExampleIndexQuery, @@ -30,14 +31,18 @@ class GremlinExampleIndexQueryNode(BaseNode): operator: GremlinExampleIndexQuery def node_init(self): + # Lazy import to avoid circular dependency + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel + # Build operator (index lazy-loading handled in operator) + vector_index = get_vector_index_class(index_settings.cur_vector_index) embedding = Embeddings().get_embedding() example_num = getattr(self.wk_input, "example_num", None) if not isinstance(example_num, int): example_num = 2 # Clamp to [0, 10] example_num = max(0, min(10, example_num)) - self.operator = GremlinExampleIndexQuery(embedding=embedding, num_examples=example_num) + self.operator = GremlinExampleIndexQuery(vector_index=vector_index, embedding=embedding, num_examples=example_num) return CStatus() def operator_schedule(self, data_json: Dict[str, Any]): diff --git a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py index 0f815b929..1d439cd69 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py @@ -81,8 +81,8 @@ def import_schema( raise ValueError("No input data / invalid schema type") return self - def example_index_query(self, num_examples): - self.operators.append(GremlinExampleIndexQuery(self.embedding, num_examples)) + def example_index_query(self, num_examples, vector_index): + self.operators.append(GremlinExampleIndexQuery(vector_index, self.embedding, num_examples)) return self def gremlin_generate_synthesize( From 0bdbf8997edec7661ef79172a0efc1e03dbd78ba Mon Sep 17 00:00:00 2001 From: Linyu <94553312+weijinglin@users.noreply.github.com> Date: Tue, 16 Sep 2025 14:54:03 +0800 Subject: [PATCH 56/71] refactor: refactor scheduler to support dynamic workflow scheduling and pipeline pooling (#48) --- hugegraph-llm/pyproject.toml | 2 + .../src/hugegraph_llm/flows/__init__.py | 16 ++ .../hugegraph_llm/flows/build_vector_index.py | 55 ++++ .../src/hugegraph_llm/flows/common.py | 45 +++ .../src/hugegraph_llm/flows/graph_extract.py | 127 +++++++++ .../src/hugegraph_llm/flows/scheduler.py | 90 ++++++ .../models/embeddings/init_embedding.py | 36 ++- .../src/hugegraph_llm/models/llms/init_llm.py | 80 +++++- .../operators/common_op/check_schema.py | 258 ++++++++++++++++-- .../operators/document_op/chunk_split.py | 59 ++++ .../operators/hugegraph_op/schema_manager.py | 88 +++++- .../operators/index_op/build_vector_index.py | 65 ++++- .../operators/llm_op/info_extract.py | 220 +++++++++++++-- .../llm_op/property_graph_extract.py | 190 +++++++++++-- .../src/hugegraph_llm/operators/util.py | 27 ++ .../src/hugegraph_llm/state/__init__.py | 16 ++ .../src/hugegraph_llm/state/ai_state.py | 81 ++++++ .../hugegraph_llm/utils/graph_index_utils.py | 83 ++++-- .../hugegraph_llm/utils/vector_index_utils.py | 67 +++-- 19 files changed, 1472 insertions(+), 133 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/__init__.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/common.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/scheduler.py create mode 100644 hugegraph-llm/src/hugegraph_llm/operators/util.py create mode 100644 hugegraph-llm/src/hugegraph_llm/state/__init__.py create mode 100644 hugegraph-llm/src/hugegraph_llm/state/ai_state.py diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index 1bd3b748c..2b0f29ace 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -61,6 +61,7 @@ dependencies = [ "apscheduler", "litellm", "hugegraph-python-client", + "pycgraph", ] [project.urls] homepage = "https://hugegraph.apache.org/" @@ -88,3 +89,4 @@ allow-direct-references = true [tool.uv.sources] hugegraph-python-client = { workspace = true } +pycgraph = { git = "https://github.com/ChunelFeng/CGraph.git", subdirectory = "python", rev = "main", marker = "sys_platform == 'linux'" } diff --git a/hugegraph-llm/src/hugegraph_llm/flows/__init__.py b/hugegraph-llm/src/hugegraph_llm/flows/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/__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/flows/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py new file mode 100644 index 000000000..f1ee8c1c4 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py @@ -0,0 +1,55 @@ +# 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.common import BaseFlow +from hugegraph_llm.state.ai_state import WkFlowInput + +import json +from PyCGraph import GPipeline + +from hugegraph_llm.operators.document_op.chunk_split import ChunkSplitNode +from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndexNode +from hugegraph_llm.state.ai_state import WkFlowState + + +class BuildVectorIndexFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, texts): + prepared_input.texts = texts + prepared_input.language = "zh" + prepared_input.split_type = "paragraph" + return + + def build_flow(self, texts): + pipeline = GPipeline() + # prepare for workflow input + prepared_input = WkFlowInput() + self.prepare(prepared_input, texts) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + chunk_split_node = ChunkSplitNode() + build_vector_node = BuildVectorIndexNode() + pipeline.registerGElement(chunk_split_node, set(), "chunk_split") + pipeline.registerGElement(build_vector_node, {chunk_split_node}, "build_vector") + + return pipeline + + def post_deal(self, pipeline=None): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + return json.dumps(res, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/common.py b/hugegraph-llm/src/hugegraph_llm/flows/common.py new file mode 100644 index 000000000..4c552626a --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/common.py @@ -0,0 +1,45 @@ +# 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 abc import ABC, abstractmethod + +from hugegraph_llm.state.ai_state import WkFlowInput + + +class BaseFlow(ABC): + """ + Base class for flows, defines three interface methods: prepare, build_flow, and post_deal. + """ + + @abstractmethod + def prepare(self, prepared_input: WkFlowInput, *args, **kwargs): + """ + Pre-processing interface. + """ + pass + + @abstractmethod + def build_flow(self, *args, **kwargs): + """ + Interface for building the flow. + """ + pass + + @abstractmethod + def post_deal(self, *args, **kwargs): + """ + Post-processing interface. + """ + pass diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py new file mode 100644 index 000000000..f1a6c5f6f --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py @@ -0,0 +1,127 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from PyCGraph import GPipeline +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.operators.common_op.check_schema import CheckSchemaNode +from hugegraph_llm.operators.document_op.chunk_split import ChunkSplitNode +from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManagerNode +from hugegraph_llm.operators.llm_op.info_extract import InfoExtractNode +from hugegraph_llm.operators.llm_op.property_graph_extract import ( + PropertyGraphExtractNode, +) +from hugegraph_llm.utils.log import log + + +class GraphExtractFlow(BaseFlow): + def __init__(self): + pass + + def _import_schema( + self, + from_hugegraph=None, + from_extraction=None, + from_user_defined=None, + ): + if from_hugegraph: + return SchemaManagerNode() + elif from_user_defined: + return CheckSchemaNode() + elif from_extraction: + raise NotImplementedError("Not implemented yet") + else: + raise ValueError("No input data / invalid schema type") + + def prepare( + self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type + ): + # prepare input data + prepared_input.texts = texts + prepared_input.language = "zh" + prepared_input.split_type = "document" + prepared_input.example_prompt = example_prompt + prepared_input.schema = schema + schema = schema.strip() + if schema.startswith("{"): + try: + schema = json.loads(schema) + prepared_input.schema = schema + except json.JSONDecodeError as exc: + log.error("Invalid JSON format in schema. Please check it again.") + raise ValueError("Invalid JSON format in schema.") from exc + else: + log.info("Get schema '%s' from graphdb.", schema) + prepared_input.graph_name = schema + return + + def build_flow(self, schema, texts, example_prompt, extract_type): + pipeline = GPipeline() + prepared_input = WkFlowInput() + # prepare input data + self.prepare(prepared_input, schema, texts, example_prompt, extract_type) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + schema = schema.strip() + schema_node = None + if schema.startswith("{"): + try: + schema = json.loads(schema) + schema_node = self._import_schema(from_user_defined=schema) + except json.JSONDecodeError as exc: + log.error("Invalid JSON format in schema. Please check it again.") + raise ValueError("Invalid JSON format in schema.") from exc + else: + log.info("Get schema '%s' from graphdb.", schema) + schema_node = self._import_schema(from_hugegraph=schema) + + chunk_split_node = ChunkSplitNode() + graph_extract_node = None + if extract_type == "triples": + graph_extract_node = InfoExtractNode() + elif extract_type == "property_graph": + graph_extract_node = PropertyGraphExtractNode() + else: + raise ValueError(f"Unsupported extract_type: {extract_type}") + pipeline.registerGElement(schema_node, set(), "schema_node") + pipeline.registerGElement(chunk_split_node, set(), "chunk_split") + pipeline.registerGElement( + graph_extract_node, {schema_node, chunk_split_node}, "graph_extract" + ) + + return pipeline + + def post_deal(self, pipeline=None): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + vertices = res.get("vertices", []) + edges = res.get("edges", []) + if not vertices and not edges: + log.info("Please check the schema.(The schema may not match the Doc)") + return json.dumps( + { + "vertices": vertices, + "edges": edges, + "warning": "The schema may not match the Doc", + }, + ensure_ascii=False, + indent=2, + ) + return json.dumps( + {"vertices": vertices, "edges": edges}, + ensure_ascii=False, + indent=2, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py new file mode 100644 index 000000000..b096310db --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -0,0 +1,90 @@ +# 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 threading +from typing import Dict, Any +from PyCGraph import GPipelineManager +from hugegraph_llm.flows.build_vector_index import BuildVectorIndexFlow +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.flows.graph_extract import GraphExtractFlow +from hugegraph_llm.utils.log import log + + +class Scheduler: + pipeline_pool: Dict[str, Any] = None + max_pipeline: int + + def __init__(self, max_pipeline: int = 10): + self.pipeline_pool = {} + # pipeline_pool act as a manager of GPipelineManager which used for pipeline management + self.pipeline_pool["build_vector_index"] = { + "manager": GPipelineManager(), + "flow": BuildVectorIndexFlow(), + } + self.pipeline_pool["graph_extract"] = { + "manager": GPipelineManager(), + "flow": GraphExtractFlow(), + } + self.max_pipeline = max_pipeline + + # TODO: Implement Agentic Workflow + def agentic_flow(self): + pass + + def schedule_flow(self, flow: str, *args, **kwargs): + if flow not in self.pipeline_pool: + raise ValueError(f"Unsupported workflow {flow}") + manager = self.pipeline_pool[flow]["manager"] + flow: BaseFlow = self.pipeline_pool[flow]["flow"] + pipeline = manager.fetch() + if pipeline is None: + # call coresponding flow_func to create new workflow + pipeline = flow.build_flow(*args, **kwargs) + status = pipeline.init() + if status.isErr(): + error_msg = f"Error in flow init: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + status = pipeline.run() + if status.isErr(): + error_msg = f"Error in flow execution: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + res = flow.post_deal(pipeline) + manager.add(pipeline) + return res + else: + # fetch pipeline & prepare input for flow + prepared_input = pipeline.getGParamWithNoEmpty("wkflow_input") + flow.prepare(prepared_input, *args, **kwargs) + status = pipeline.run() + if status.isErr(): + raise RuntimeError(f"Error in flow execution {status.getInfo()}") + res = flow.post_deal(pipeline) + manager.release(pipeline) + return res + + +class SchedulerSingleton: + _instance = None + _instance_lock = threading.Lock() + + @classmethod + def get_instance(cls): + if cls._instance is None: + with cls._instance_lock: + if cls._instance is None: + cls._instance = Scheduler() + return cls._instance diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index 48e4968c4..3ad50b3ec 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -17,10 +17,40 @@ from hugegraph_llm.config import llm_settings +from hugegraph_llm.config import LLMConfig from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding from hugegraph_llm.models.embeddings.ollama import OllamaEmbedding from hugegraph_llm.models.embeddings.openai import OpenAIEmbedding +model_map = { + "openai": llm_settings.openai_embedding_model, + "ollama/local": llm_settings.ollama_embedding_model, + "litellm": llm_settings.litellm_embedding_model, +} + + +def get_embedding(llm_settings: LLMConfig): + if llm_settings.embedding_type == "openai": + return OpenAIEmbedding( + model_name=llm_settings.openai_embedding_model, + api_key=llm_settings.openai_embedding_api_key, + api_base=llm_settings.openai_embedding_api_base, + ) + if llm_settings.embedding_type == "ollama/local": + return OllamaEmbedding( + model_name=llm_settings.ollama_embedding_model, + host=llm_settings.ollama_embedding_host, + port=llm_settings.ollama_embedding_port, + ) + if llm_settings.embedding_type == "litellm": + return LiteLLMEmbedding( + model_name=llm_settings.litellm_embedding_model, + api_key=llm_settings.litellm_embedding_api_key, + api_base=llm_settings.litellm_embedding_api_base, + ) + + raise Exception("embedding type is not supported !") + class Embeddings: def __init__(self): @@ -31,19 +61,19 @@ def get_embedding(self): return OpenAIEmbedding( model_name=llm_settings.openai_embedding_model, api_key=llm_settings.openai_embedding_api_key, - api_base=llm_settings.openai_embedding_api_base + api_base=llm_settings.openai_embedding_api_base, ) if self.embedding_type == "ollama/local": return OllamaEmbedding( model_name=llm_settings.ollama_embedding_model, host=llm_settings.ollama_embedding_host, - port=llm_settings.ollama_embedding_port + port=llm_settings.ollama_embedding_port, ) if self.embedding_type == "litellm": return LiteLLMEmbedding( model_name=llm_settings.litellm_embedding_model, api_key=llm_settings.litellm_embedding_api_key, - api_base=llm_settings.litellm_embedding_api_base + api_base=llm_settings.litellm_embedding_api_base, ) raise Exception("embedding type is not supported !") 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 e70b0d9d7..7e1eaab68 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py @@ -15,13 +15,85 @@ # specific language governing permissions and limitations # under the License. - +from hugegraph_llm.config import LLMConfig from hugegraph_llm.models.llms.ollama import OllamaClient from hugegraph_llm.models.llms.openai import OpenAIClient from hugegraph_llm.models.llms.litellm import LiteLLMClient from hugegraph_llm.config import llm_settings +def get_chat_llm(llm_settings: LLMConfig): + if llm_settings.chat_llm_type == "openai": + return OpenAIClient( + api_key=llm_settings.openai_chat_api_key, + api_base=llm_settings.openai_chat_api_base, + model_name=llm_settings.openai_chat_language_model, + max_tokens=llm_settings.openai_chat_tokens, + ) + if llm_settings.chat_llm_type == "ollama/local": + return OllamaClient( + model=llm_settings.ollama_chat_language_model, + host=llm_settings.ollama_chat_host, + port=llm_settings.ollama_chat_port, + ) + if llm_settings.chat_llm_type == "litellm": + return LiteLLMClient( + api_key=llm_settings.litellm_chat_api_key, + api_base=llm_settings.litellm_chat_api_base, + model_name=llm_settings.litellm_chat_language_model, + max_tokens=llm_settings.litellm_chat_tokens, + ) + raise Exception("chat llm type is not supported !") + + +def get_extract_llm(llm_settings: LLMConfig): + if llm_settings.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, + ) + if llm_settings.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, + ) + if llm_settings.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, + ) + raise Exception("extract llm type is not supported !") + + +def get_text2gql_llm(llm_settings: LLMConfig): + if llm_settings.text2gql_llm_type == "openai": + return OpenAIClient( + api_key=llm_settings.openai_text2gql_api_key, + api_base=llm_settings.openai_text2gql_api_base, + model_name=llm_settings.openai_text2gql_language_model, + max_tokens=llm_settings.openai_text2gql_tokens, + ) + if llm_settings.text2gql_llm_type == "ollama/local": + return OllamaClient( + model=llm_settings.ollama_text2gql_language_model, + host=llm_settings.ollama_text2gql_host, + port=llm_settings.ollama_text2gql_port, + ) + if llm_settings.text2gql_llm_type == "litellm": + return LiteLLMClient( + api_key=llm_settings.litellm_text2gql_api_key, + api_base=llm_settings.litellm_text2gql_api_base, + model_name=llm_settings.litellm_text2gql_language_model, + max_tokens=llm_settings.litellm_text2gql_tokens, + ) + raise Exception("text2gql llm type is not supported !") + + class LLMs: def __init__(self): self.chat_llm_type = llm_settings.chat_llm_type @@ -101,4 +173,8 @@ def get_text2gql_llm(self): if __name__ == "__main__": client = LLMs().get_chat_llm() print(client.generate(prompt="What is the capital of China?")) - print(client.generate(messages=[{"role": "user", "content": "What is the capital of China?"}])) + print( + client.generate( + messages=[{"role": "user", "content": "What is the capital of China?"}] + ) + ) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index 3220d9f3d..7a533517a 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -20,8 +20,12 @@ from hugegraph_llm.enums.property_cardinality import PropertyCardinality from hugegraph_llm.enums.property_data_type import PropertyDataType +from hugegraph_llm.operators.util import init_context from hugegraph_llm.utils.log import log +from PyCGraph import GNode, CStatus +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + def log_and_raise(message: str) -> None: log.warning(message) @@ -59,64 +63,270 @@ def _validate_schema(self, schema: Dict[str, Any]) -> None: check_type(schema, dict, "Input data is not a dictionary.") if "vertexlabels" not in schema or "edgelabels" not in schema: log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") - check_type(schema["vertexlabels"], list, "'vertexlabels' in input data is not a list.") - check_type(schema["edgelabels"], list, "'edgelabels' in input data is not a list.") + check_type( + schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." + ) + check_type( + schema["edgelabels"], list, "'edgelabels' in input data is not a list." + ) + + def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): + property_labels = schema.get("propertykeys", []) + check_type( + property_labels, + list, + "'propertykeys' in input data is not of correct type.", + ) + property_label_set = {label["name"] for label in property_labels} + return property_labels, property_label_set + + def _process_vertex_labels( + self, schema: Dict[str, Any], property_labels: list, property_label_set: set + ) -> None: + for vertex_label in schema["vertexlabels"]: + self._validate_vertex_label(vertex_label) + properties = vertex_label["properties"] + primary_keys = self._process_keys( + vertex_label, "primary_keys", properties[:1] + ) + if len(primary_keys) == 0: + log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") + vertex_label["primary_keys"] = primary_keys + nullable_keys = self._process_keys( + vertex_label, "nullable_keys", properties[1:] + ) + vertex_label["nullable_keys"] = nullable_keys + self._add_missing_properties( + properties, property_labels, property_label_set + ) + + def _process_edge_labels( + self, schema: Dict[str, Any], property_labels: list, property_label_set: set + ) -> None: + for edge_label in schema["edgelabels"]: + self._validate_edge_label(edge_label) + properties = edge_label.get("properties", []) + self._add_missing_properties( + properties, property_labels, property_label_set + ) + + def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: + check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") + if "name" not in vertex_label: + log_and_raise("VertexLabel in input data does not contain 'name'.") + check_type( + vertex_label["name"], str, "'name' in vertex_label is not of correct type." + ) + if "properties" not in vertex_label: + log_and_raise("VertexLabel in input data does not contain 'properties'.") + check_type( + vertex_label["properties"], + list, + "'properties' in vertex_label is not of correct type.", + ) + if len(vertex_label["properties"]) == 0: + log_and_raise("'properties' in vertex_label is empty.") + + def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: + check_type(edge_label, dict, "EdgeLabel in input data is not a dictionary.") + if ( + "name" not in edge_label + or "source_label" not in edge_label + or "target_label" not in edge_label + ): + log_and_raise( + "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." + ) + check_type( + edge_label["name"], str, "'name' in edge_label is not of correct type." + ) + check_type( + edge_label["source_label"], + str, + "'source_label' in edge_label is not of correct type.", + ) + check_type( + edge_label["target_label"], + str, + "'target_label' in edge_label is not of correct type.", + ) + + def _process_keys( + self, label: Dict[str, Any], key_type: str, default_keys: list + ) -> list: + keys = label.get(key_type, default_keys) + check_type( + keys, list, f"'{key_type}' in {label['name']} is not of correct type." + ) + new_keys = [key for key in keys if key in label["properties"]] + return new_keys + + def _add_missing_properties( + self, properties: list, property_labels: list, property_label_set: set + ) -> None: + for prop in properties: + if prop not in property_label_set: + property_labels.append( + { + "name": prop, + "data_type": PropertyDataType.DEFAULT.value, + "cardinality": PropertyCardinality.DEFAULT.value, + } + ) + property_label_set.add(prop) + + +class CheckSchemaNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + return init_context(self) + + def node_init(self): + if self.wk_input.schema is None: + return CStatus(-1, "Error occurs when prepare for workflow input") + self.data = self.wk_input.schema + return CStatus() + + def run(self) -> CStatus: + # init workflow input + sts = self.node_init() + if sts.isErr(): + return sts + # 1. Validate the schema structure + self.context.lock() + schema = self.data or self.context.schema + self._validate_schema(schema) + # 2. Process property labels and also create a set for it + property_labels, property_label_set = self._process_property_labels(schema) + # 3. Process properties in given vertex/edge labels + self._process_vertex_labels(schema, property_labels, property_label_set) + self._process_edge_labels(schema, property_labels, property_label_set) + # 4. Update schema with processed pks + schema["propertykeys"] = property_labels + self.context.schema = schema + self.context.unlock() + return CStatus() + + def _validate_schema(self, schema: Dict[str, Any]) -> None: + check_type(schema, dict, "Input data is not a dictionary.") + if "vertexlabels" not in schema or "edgelabels" not in schema: + log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") + check_type( + schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." + ) + check_type( + schema["edgelabels"], list, "'edgelabels' in input data is not a list." + ) def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): property_labels = schema.get("propertykeys", []) - check_type(property_labels, list, "'propertykeys' in input data is not of correct type.") + check_type( + property_labels, + list, + "'propertykeys' in input data is not of correct type.", + ) property_label_set = {label["name"] for label in property_labels} return property_labels, property_label_set - def _process_vertex_labels(self, schema: Dict[str, Any], property_labels: list, property_label_set: set) -> None: + def _process_vertex_labels( + self, schema: Dict[str, Any], property_labels: list, property_label_set: set + ) -> None: for vertex_label in schema["vertexlabels"]: self._validate_vertex_label(vertex_label) properties = vertex_label["properties"] - primary_keys = self._process_keys(vertex_label, "primary_keys", properties[:1]) + primary_keys = self._process_keys( + vertex_label, "primary_keys", properties[:1] + ) if len(primary_keys) == 0: log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") vertex_label["primary_keys"] = primary_keys - nullable_keys = self._process_keys(vertex_label, "nullable_keys", properties[1:]) + nullable_keys = self._process_keys( + vertex_label, "nullable_keys", properties[1:] + ) vertex_label["nullable_keys"] = nullable_keys - self._add_missing_properties(properties, property_labels, property_label_set) + self._add_missing_properties( + properties, property_labels, property_label_set + ) - def _process_edge_labels(self, schema: Dict[str, Any], property_labels: list, property_label_set: set) -> None: + def _process_edge_labels( + self, schema: Dict[str, Any], property_labels: list, property_label_set: set + ) -> None: for edge_label in schema["edgelabels"]: self._validate_edge_label(edge_label) properties = edge_label.get("properties", []) - self._add_missing_properties(properties, property_labels, property_label_set) + self._add_missing_properties( + properties, property_labels, property_label_set + ) def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") if "name" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'name'.") - check_type(vertex_label["name"], str, "'name' in vertex_label is not of correct type.") + check_type( + vertex_label["name"], str, "'name' in vertex_label is not of correct type." + ) if "properties" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'properties'.") - check_type(vertex_label["properties"], list, "'properties' in vertex_label is not of correct type.") + check_type( + vertex_label["properties"], + list, + "'properties' in vertex_label is not of correct type.", + ) if len(vertex_label["properties"]) == 0: log_and_raise("'properties' in vertex_label is empty.") def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: check_type(edge_label, dict, "EdgeLabel in input data is not a dictionary.") - if "name" not in edge_label or "source_label" not in edge_label or "target_label" not in edge_label: - log_and_raise("EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'.") - check_type(edge_label["name"], str, "'name' in edge_label is not of correct type.") - check_type(edge_label["source_label"], str, "'source_label' in edge_label is not of correct type.") - check_type(edge_label["target_label"], str, "'target_label' in edge_label is not of correct type.") + if ( + "name" not in edge_label + or "source_label" not in edge_label + or "target_label" not in edge_label + ): + log_and_raise( + "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." + ) + check_type( + edge_label["name"], str, "'name' in edge_label is not of correct type." + ) + check_type( + edge_label["source_label"], + str, + "'source_label' in edge_label is not of correct type.", + ) + check_type( + edge_label["target_label"], + str, + "'target_label' in edge_label is not of correct type.", + ) - def _process_keys(self, label: Dict[str, Any], key_type: str, default_keys: list) -> list: + def _process_keys( + self, label: Dict[str, Any], key_type: str, default_keys: list + ) -> list: keys = label.get(key_type, default_keys) - check_type(keys, list, f"'{key_type}' in {label['name']} is not of correct type.") + check_type( + keys, list, f"'{key_type}' in {label['name']} is not of correct type." + ) new_keys = [key for key in keys if key in label["properties"]] return new_keys - def _add_missing_properties(self, properties: list, property_labels: list, property_label_set: set) -> None: + def _add_missing_properties( + self, properties: list, property_labels: list, property_label_set: set + ) -> None: for prop in properties: if prop not in property_label_set: - property_labels.append({ - "name": prop, - "data_type": PropertyDataType.DEFAULT.value, - "cardinality": PropertyCardinality.DEFAULT.value, - }) + property_labels.append( + { + "name": prop, + "data_type": PropertyDataType.DEFAULT.value, + "cardinality": PropertyCardinality.DEFAULT.value, + } + ) property_label_set.add(prop) + + def get_result(self): + self.context.lock() + res = self.context.to_json() + self.context.unlock() + return res 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 8c2dd80f5..d779a40ab 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 @@ -19,6 +19,8 @@ from typing import Literal, Dict, Any, Optional, Union, List from langchain_text_splitters import RecursiveCharacterTextSplitter +from hugegraph_llm.operators.util import init_context +from PyCGraph import GNode, CStatus # Constants LANGUAGE_ZH = "zh" @@ -27,6 +29,63 @@ SPLIT_TYPE_PARAGRAPH = "paragraph" SPLIT_TYPE_SENTENCE = "sentence" + +class ChunkSplitNode(GNode): + def init(self): + return init_context(self) + + def node_init(self): + if ( + self.wk_input.texts is None + or self.wk_input.language is None + or self.wk_input.split_type is None + ): + return CStatus(-1, "Error occurs when prepare for workflow input") + texts = self.wk_input.texts + language = self.wk_input.language + split_type = self.wk_input.split_type + if isinstance(texts, str): + texts = [texts] + self.texts = texts + self.separators = self._get_separators(language) + self.text_splitter = self._get_text_splitter(split_type) + return CStatus() + + def _get_separators(self, language: str) -> List[str]: + if language == LANGUAGE_ZH: + return ["\n\n", "\n", "。", "，", ""] + if language == LANGUAGE_EN: + return ["\n\n", "\n", ".", ",", " ", ""] + raise ValueError("language must be zh or en") + + 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 + if split_type == SPLIT_TYPE_SENTENCE: + return RecursiveCharacterTextSplitter( + chunk_size=50, chunk_overlap=0, separators=self.separators + ).split_text + raise ValueError("Type must be document, paragraph or sentence") + + def run(self): + sts = self.node_init() + if sts.isErr(): + return sts + all_chunks = [] + for text in self.texts: + chunks = self.text_splitter(text) + all_chunks.extend(chunks) + + self.context.lock() + self.context.chunks = all_chunks + self.context.unlock() + return CStatus() + + class ChunkSplit: def __init__( self, 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 2f50bb818..670c18b4a 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 @@ -17,8 +17,12 @@ from typing import Dict, Any, Optional from hugegraph_llm.config import huge_settings +from hugegraph_llm.operators.util import init_context +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState from pyhugegraph.client import PyHugeClient +from PyCGraph import GNode, CStatus + class SchemaManager: def __init__(self, graph_name: str): @@ -39,15 +43,22 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} + new_vertex = { + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex + } mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) if "edgelabels" in schema: mini_schema["edgelabels"] = [] for edge in schema["edgelabels"]: - new_edge = {key: edge[key] for key in - ["name", "source_label", "target_label", "properties"] if key in edge} + new_edge = { + key: edge[key] + for key in ["name", "source_label", "target_label", "properties"] + if key in edge + } mini_schema["edgelabels"].append(new_edge) return mini_schema @@ -63,3 +74,74 @@ def run(self, context: Optional[Dict[str, Any]]) -> Dict[str, Any]: # TODO: enhance the logic here context["simple_schema"] = self.simple_schema(schema) return context + + +class SchemaManagerNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + return init_context(self) + + def node_init(self): + if self.wk_input.graph_name is None: + return CStatus(-1, "Error occurs when prepare for workflow input") + graph_name = self.wk_input.graph_name + self.graph_name = graph_name + self.client = PyHugeClient( + url=huge_settings.graph_url, + graph=self.graph_name, + user=huge_settings.graph_user, + pwd=huge_settings.graph_pwd, + graphspace=huge_settings.graph_space, + ) + self.schema = self.client.schema() + return CStatus() + + def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: + mini_schema = {} + + # Add necessary vertexlabels items (3) + if "vertexlabels" in schema: + mini_schema["vertexlabels"] = [] + for vertex in schema["vertexlabels"]: + new_vertex = { + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex + } + mini_schema["vertexlabels"].append(new_vertex) + + # Add necessary edgelabels items (4) + if "edgelabels" in schema: + mini_schema["edgelabels"] = [] + for edge in schema["edgelabels"]: + new_edge = { + key: edge[key] + for key in ["name", "source_label", "target_label", "properties"] + if key in edge + } + mini_schema["edgelabels"].append(new_edge) + + return mini_schema + + def run(self) -> CStatus: + sts = self.node_init() + if sts.isErr(): + return sts + schema = self.schema.getSchema() + if not schema["vertexlabels"] and not schema["edgelabels"]: + raise Exception(f"Can not get {self.graph_name}'s schema from HugeGraph!") + + self.context.lock() + self.context.schema = schema + # TODO: enhance the logic here + self.context.simple_schema = self.simple_schema(schema) + self.context.unlock() + return CStatus() + + def get_result(self): + self.context.lock() + res = self.context.to_json() + self.context.unlock() + return res diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py index ffb35564b..ee89d330f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py @@ -23,20 +23,75 @@ from hugegraph_llm.config import huge_settings, resource_path, llm_settings from hugegraph_llm.indices.vector_index import VectorIndex from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel, get_filename_prefix, get_index_folder_name +from hugegraph_llm.utils.embedding_utils import ( + get_embeddings_parallel, + get_filename_prefix, + get_index_folder_name, +) from hugegraph_llm.utils.log import log +from hugegraph_llm.operators.util import init_context +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from PyCGraph import GNode, CStatus + + +class BuildVectorIndexNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + return init_context(self) + + def node_init(self): + self.embedding = get_embedding(llm_settings) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(self.embedding, "model_name", None) + ) + self.vector_index = VectorIndex.from_index_file( + self.index_dir, self.filename_prefix + ) + return CStatus() + + def run(self): + # init workflow input + sts = self.node_init() + if sts.isErr(): + return sts + self.context.lock() + try: + if self.context.chunks is None: + raise ValueError("chunks not found in context.") + chunks = self.context.chunks + finally: + self.context.unlock() + chunks_embedding = [] + log.debug("Building vector index for %s chunks...", len(chunks)) + # TODO: use async_get_texts_embedding instead of single sync method + chunks_embedding = asyncio.run(get_embeddings_parallel(self.embedding, chunks)) + if len(chunks_embedding) > 0: + self.vector_index.add(chunks_embedding, chunks) + self.vector_index.to_index_file(self.index_dir, self.filename_prefix) + return CStatus() + class BuildVectorIndex: def __init__(self, embedding: BaseEmbedding): self.embedding = embedding - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, - getattr(self.embedding, "model_name", None) + llm_settings.embedding_type, getattr(self.embedding, "model_name", None) + ) + self.vector_index = VectorIndex.from_index_file( + self.index_dir, self.filename_prefix ) - self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: if "chunks" not in context: 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 42bb6b108..15a8fdda7 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 @@ -18,20 +18,26 @@ import re from typing import List, Any, Dict, Optional +from hugegraph_llm.config import llm_settings from hugegraph_llm.document.chunk_split import ChunkSplitter from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.utils.log import log +from hugegraph_llm.operators.util import init_context +from hugegraph_llm.models.llms.init_llm import get_chat_llm +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from PyCGraph import GNode, CStatus + SCHEMA_EXAMPLE_PROMPT = """## Main Task Extract Triples from the given text and graph schema ## Basic Rules 1. The output format must be: (X,Y,Z) - LABEL -In this format, Y must be a value from "properties" or "edge_label", +In this format, Y must be a value from "properties" or "edge_label", and LABEL must be X's vertex_label or Y's edge_label. 2. Don't extract attribute/property fields that do not exist in the given schema 3. Ensure the extract property is in the same type as the schema (like 'age' should be a number) -4. Translate the given schema filed into Chinese if the given text is Chinese but the schema is in English (Optional) +4. Translate the given schema filed into Chinese if the given text is Chinese but the schema is in English (Optional) ## Example (Note: Update the example to correspond to the given text and schema) ### Input example: @@ -75,8 +81,10 @@ def generate_extract_triple_prompt(text, schema=None) -> str: if schema: return schema_real_prompt - log.warning("Recommend to provide a graph schema to improve the extraction accuracy. " - "Now using the default schema.") + log.warning( + "Recommend to provide a graph schema to improve the extraction accuracy. " + "Now using the default schema." + ) return text_based_prompt @@ -105,11 +113,17 @@ def extract_triples_by_regex_with_schema(schema, text, graph): # TODO: use a more efficient way to compare the extract & input property p_lower = p.lower() for vertex in schema["vertices"]: - if vertex["vertex_label"] == label and any(pp.lower() == p_lower - for pp in vertex["properties"]): + if vertex["vertex_label"] == label and any( + pp.lower() == p_lower for pp in vertex["properties"] + ): id = f"{label}-{s}" if id not in vertices_dict: - vertices_dict[id] = {"id": id, "name": s, "label": label, "properties": {p: o}} + vertices_dict[id] = { + "id": id, + "name": s, + "label": label, + "properties": {p: o}, + } else: vertices_dict[id]["properties"].update({p: o}) break @@ -118,25 +132,35 @@ def extract_triples_by_regex_with_schema(schema, text, graph): source_label = edge["source_vertex_label"] source_id = f"{source_label}-{s}" if source_id not in vertices_dict: - vertices_dict[source_id] = {"id": source_id, "name": s, "label": source_label, - "properties": {}} + vertices_dict[source_id] = { + "id": source_id, + "name": s, + "label": source_label, + "properties": {}, + } target_label = edge["target_vertex_label"] target_id = f"{target_label}-{o}" if target_id not in vertices_dict: - vertices_dict[target_id] = {"id": target_id, "name": o, "label": target_label, - "properties": {}} - graph["edges"].append({"start": source_id, "end": target_id, "type": label, - "properties": {}}) + vertices_dict[target_id] = { + "id": target_id, + "name": o, + "label": target_label, + "properties": {}, + } + graph["edges"].append( + { + "start": source_id, + "end": target_id, + "type": label, + "properties": {}, + } + ) break - graph["vertices"] = vertices_dict.values() + graph["vertices"] = list(vertices_dict.values()) class InfoExtract: - def __init__( - self, - llm: BaseLLM, - example_prompt: Optional[str] = None - ) -> None: + def __init__(self, llm: BaseLLM, example_prompt: Optional[str] = None) -> None: self.llm = llm self.example_prompt = example_prompt @@ -152,7 +176,12 @@ def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]: for sentence in chunks: proceeded_chunk = self.extract_triples_by_llm(schema, sentence) - log.debug("[Legacy] %s input: %s \n output:%s", self.__class__.__name__, sentence, proceeded_chunk) + log.debug( + "[Legacy] %s input: %s \n output:%s", + self.__class__.__name__, + sentence, + proceeded_chunk, + ) if schema: extract_triples_by_regex_with_schema(schema, proceeded_chunk, context) else: @@ -175,7 +204,152 @@ def valid(self, element_id: str, max_length: int = 256) -> bool: return True def _filter_long_id(self, graph) -> Dict[str, List[Any]]: - graph["vertices"] = [vertex for vertex in graph["vertices"] if self.valid(vertex["id"])] - graph["edges"] = [edge for edge in graph["edges"] - if self.valid(edge["start"]) and self.valid(edge["end"])] + graph["vertices"] = [ + vertex for vertex in graph["vertices"] if self.valid(vertex["id"]) + ] + graph["edges"] = [ + edge + for edge in graph["edges"] + if self.valid(edge["start"]) and self.valid(edge["end"]) + ] return graph + + +class InfoExtractNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + return init_context(self) + + def node_init(self): + self.llm = get_chat_llm(llm_settings) + if self.wk_input.example_prompt is None: + return CStatus(-1, "Error occurs when prepare for workflow input") + self.example_prompt = self.wk_input.example_prompt + return CStatus() + + def extract_triples_by_regex_with_schema(self, schema, text): + text = text.replace("\\n", " ").replace("\\", " ").replace("\n", " ") + pattern = r"\((.*?), (.*?), (.*?)\) - ([^ ]*)" + matches = re.findall(pattern, text) + + vertices_dict = {v["id"]: v for v in self.context.vertices} + for match in matches: + s, p, o, label = [item.strip() for item in match] + if None in [label, s, p, o]: + continue + # TODO: use a more efficient way to compare the extract & input property + p_lower = p.lower() + for vertex in schema["vertices"]: + if vertex["vertex_label"] == label and any( + pp.lower() == p_lower for pp in vertex["properties"] + ): + id = f"{label}-{s}" + if id not in vertices_dict: + vertices_dict[id] = { + "id": id, + "name": s, + "label": label, + "properties": {p: o}, + } + else: + vertices_dict[id]["properties"].update({p: o}) + break + for edge in schema["edges"]: + if edge["edge_label"] == label: + source_label = edge["source_vertex_label"] + source_id = f"{source_label}-{s}" + if source_id not in vertices_dict: + vertices_dict[source_id] = { + "id": source_id, + "name": s, + "label": source_label, + "properties": {}, + } + target_label = edge["target_vertex_label"] + target_id = f"{target_label}-{o}" + if target_id not in vertices_dict: + vertices_dict[target_id] = { + "id": target_id, + "name": o, + "label": target_label, + "properties": {}, + } + self.context.edges.append( + { + "start": source_id, + "end": target_id, + "type": label, + "properties": {}, + } + ) + break + self.context.vertices = list(vertices_dict.values()) + + def extract_triples_by_regex(self, text): + text = text.replace("\\n", " ").replace("\\", " ").replace("\n", " ") + pattern = r"\((.*?), (.*?), (.*?)\)" + self.context.triples += re.findall(pattern, text) + + def run(self) -> CStatus: + sts = self.node_init() + if sts.isErr(): + return sts + self.context.lock() + if self.context.chunks is None: + self.context.unlock() + raise ValueError("parameter required by extract node not found in context.") + schema = self.context.schema + chunks = self.context.chunks + + if schema: + self.context.vertices = [] + self.context.edges = [] + else: + self.context.triples = [] + + self.context.unlock() + + for sentence in chunks: + proceeded_chunk = self.extract_triples_by_llm(schema, sentence) + log.debug( + "[Legacy] %s input: %s \n output:%s", + self.__class__.__name__, + sentence, + proceeded_chunk, + ) + if schema: + self.extract_triples_by_regex_with_schema(schema, proceeded_chunk) + else: + self.extract_triples_by_regex(proceeded_chunk) + + if self.context.call_count: + self.context.call_count += len(chunks) + else: + self.context.call_count = len(chunks) + self._filter_long_id() + return CStatus() + + def extract_triples_by_llm(self, schema, chunk) -> str: + prompt = generate_extract_triple_prompt(chunk, schema) + if self.example_prompt is not None: + prompt = self.example_prompt + prompt + return self.llm.generate(prompt=prompt) + + # TODO: make 'max_length' be a configurable param in settings.py/settings.cfg + def valid(self, element_id: str, max_length: int = 256) -> bool: + if len(element_id.encode("utf-8")) >= max_length: + log.warning("Filter out GraphElementID too long: %s", element_id) + return False + return True + + def _filter_long_id(self): + self.context.vertices = [ + vertex for vertex in self.context.vertices if self.valid(vertex["id"]) + ] + self.context.edges = [ + edge + for edge in self.context.edges + if self.valid(edge["start"]) and self.valid(edge["end"]) + ] 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 faff1c6b2..6e492b8f5 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 @@ -21,16 +21,19 @@ import re from typing import List, Any, Dict -from hugegraph_llm.config import prompt +from hugegraph_llm.config import llm_settings, prompt from hugegraph_llm.document.chunk_split import ChunkSplitter from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.utils.log import log -""" -TODO: It is not clear whether there is any other dependence on the SCHEMA_EXAMPLE_PROMPT variable. -Because the SCHEMA_EXAMPLE_PROMPT variable will no longer change based on -prompt.extract_graph_prompt changes after the system loads, this does not seem to meet expectations. -""" +from hugegraph_llm.operators.util import init_context +from hugegraph_llm.models.llms.init_llm import get_chat_llm +from hugegraph_llm.state.ai_state import WkFlowState, WkFlowInput +from PyCGraph import GNode, CStatus + +# TODO: It is not clear whether there is any other dependence on the SCHEMA_EXAMPLE_PROMPT variable. +# Because the SCHEMA_EXAMPLE_PROMPT variable will no longer change based on +# prompt.extract_graph_prompt changes after the system loads, this does not seem to meet expectations. SCHEMA_EXAMPLE_PROMPT = prompt.extract_graph_prompt @@ -60,20 +63,18 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: properties_map["vertex"][vertex["name"]] = { "primary_keys": vertex["primary_keys"], "nullable_keys": vertex["nullable_keys"], - "properties": vertex["properties"] + "properties": vertex["properties"], } for edge in schema["edgelabels"]: - properties_map["edge"][edge["name"]] = { - "properties": edge["properties"] - } + properties_map["edge"][edge["name"]] = {"properties": edge["properties"]} log.info("properties_map: %s", properties_map) for item in items: item_type = item["type"] if item_type == "vertex": label = item["label"] - non_nullable_keys = ( - set(properties_map[item_type][label]["properties"]) - .difference(set(properties_map[item_type][label]["nullable_keys"]))) + non_nullable_keys = set( + properties_map[item_type][label]["properties"] + ).difference(set(properties_map[item_type][label]["nullable_keys"])) for key in non_nullable_keys: if key not in item["properties"]: item["properties"][key] = "NULL" @@ -87,9 +88,7 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: class PropertyGraphExtract: def __init__( - self, - llm: BaseLLM, - example_prompt: str = prompt.extract_graph_prompt + self, llm: BaseLLM, example_prompt: str = prompt.extract_graph_prompt ) -> None: self.llm = llm self.example_prompt = example_prompt @@ -105,7 +104,12 @@ def run(self, context: Dict[str, Any]) -> Dict[str, List[Any]]: items = [] for chunk in chunks: proceeded_chunk = self.extract_property_graph_by_llm(schema, chunk) - log.debug("[LLM] %s input: %s \n output:%s", self.__class__.__name__, chunk, proceeded_chunk) + log.debug( + "[LLM] %s input: %s \n output:%s", + self.__class__.__name__, + chunk, + proceeded_chunk, + ) items.extend(self._extract_and_filter_label(schema, proceeded_chunk)) items = filter_item(schema, items) for item in items: @@ -125,10 +129,132 @@ def extract_property_graph_by_llm(self, schema, chunk): def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: # Use regex to extract a JSON object with curly braces - json_match = re.search(r'({.*})', text, re.DOTALL) + json_match = re.search(r"({.*})", text, re.DOTALL) + if not json_match: + log.critical( + "Invalid property graph! No JSON object found, " + "please check the output format example in prompt." + ) + return [] + json_str = json_match.group(1).strip() + + items = [] + try: + property_graph = json.loads(json_str) + # 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'." + ) + return items + + # Create sets for valid vertex and edge labels based on the schema + vertex_label_set = {vertex["name"] for vertex in schema["vertexlabels"]} + edge_label_set = {edge["name"] for edge in schema["edgelabels"]} + + def process_items(item_list, valid_labels, item_type): + for item in item_list: + if not isinstance(item, dict): + log.warning( + "Invalid property graph item type '%s'.", type(item) + ) + continue + if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): + log.warning("Invalid item keys '%s'.", item.keys()) + continue + if item["label"] not in valid_labels: + log.warning( + "Invalid %s label '%s' has been ignored.", + item_type, + item["label"], + ) + continue + items.append(item) + + process_items(property_graph["vertices"], vertex_label_set, "vertex") + process_items(property_graph["edges"], edge_label_set, "edge") + except json.JSONDecodeError: + log.critical( + "Invalid property graph JSON! Please check the extracted JSON data carefully" + ) + return items + + +class PropertyGraphExtractNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name + return init_context(self) + + def node_init(self): + self.llm = get_chat_llm(llm_settings) + if self.wk_input.example_prompt is None: + return CStatus(-1, "Error occurs when prepare for workflow input") + self.example_prompt = self.wk_input.example_prompt + return CStatus() + + def run(self) -> CStatus: + sts = self.node_init() + if sts.isErr(): + return sts + self.context.lock() + try: + if self.context.schema is None or self.context.chunks is None: + raise ValueError( + "parameter required by extract node not found in context." + ) + schema = self.context.schema + chunks = self.context.chunks + if self.context.vertices is None: + self.context.vertices = [] + if self.context.edges is None: + self.context.edges = [] + finally: + self.context.unlock() + + items = [] + for chunk in chunks: + proceeded_chunk = self.extract_property_graph_by_llm(schema, chunk) + log.debug( + "[LLM] %s input: %s \n output:%s", + self.__class__.__name__, + chunk, + proceeded_chunk, + ) + items.extend(self._extract_and_filter_label(schema, proceeded_chunk)) + items = filter_item(schema, items) + self.context.lock() + try: + for item in items: + if item["type"] == "vertex": + self.context.vertices.append(item) + elif item["type"] == "edge": + self.context.edges.append(item) + finally: + self.context.unlock() + self.context.call_count = (self.context.call_count or 0) + len(chunks) + return CStatus() + + def extract_property_graph_by_llm(self, schema, chunk): + prompt = generate_extract_property_graph_prompt(chunk, schema) + if self.example_prompt is not None: + prompt = self.example_prompt + prompt + return self.llm.generate(prompt=prompt) + + def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: + # Use regex to extract a JSON object with curly braces + json_match = re.search(r"({.*})", text, re.DOTALL) if not json_match: - log.critical("Invalid property graph! No JSON object found, " - "please check the output format example in prompt.") + log.critical( + "Invalid property graph! No JSON object found, " + "please check the output format example in prompt." + ) return [] json_str = json_match.group(1).strip() @@ -136,8 +262,14 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: try: property_graph = json.loads(json_str) # 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 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'." + ) return items # Create sets for valid vertex and edge labels based on the schema @@ -147,18 +279,26 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: def process_items(item_list, valid_labels, item_type): for item in item_list: if not isinstance(item, dict): - log.warning("Invalid property graph item type '%s'.", type(item)) + log.warning( + "Invalid property graph item type '%s'.", type(item) + ) continue if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): log.warning("Invalid item keys '%s'.", item.keys()) continue if item["label"] not in valid_labels: - log.warning("Invalid %s label '%s' has been ignored.", item_type, item["label"]) + log.warning( + "Invalid %s label '%s' has been ignored.", + item_type, + item["label"], + ) continue items.append(item) process_items(property_graph["vertices"], vertex_label_set, "vertex") process_items(property_graph["edges"], edge_label_set, "edge") except json.JSONDecodeError: - log.critical("Invalid property graph JSON! Please check the extracted JSON data carefully") + log.critical( + "Invalid property graph JSON! Please check the extracted JSON data carefully" + ) return items diff --git a/hugegraph-llm/src/hugegraph_llm/operators/util.py b/hugegraph-llm/src/hugegraph_llm/operators/util.py new file mode 100644 index 000000000..60bdc2e86 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/operators/util.py @@ -0,0 +1,27 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from PyCGraph import CStatus + + +def init_context(obj) -> CStatus: + try: + obj.context = obj.getGParamWithNoEmpty("wkflow_state") + obj.wk_input = obj.getGParamWithNoEmpty("wkflow_input") + if obj.context is None or obj.wk_input is None: + return CStatus(-1, "Required workflow parameters not found") + return CStatus() + except Exception as e: + return CStatus(-1, f"Failed to initialize context: {str(e)}") diff --git a/hugegraph-llm/src/hugegraph_llm/state/__init__.py b/hugegraph-llm/src/hugegraph_llm/state/__init__.py new file mode 100644 index 000000000..13a83393a --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/state/__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/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py new file mode 100644 index 000000000..0543aa2b4 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -0,0 +1,81 @@ +# 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 PyCGraph import GParam, CStatus + +from typing import Union, List, Optional, Any + + +class WkFlowInput(GParam): + texts: Union[str, List[str]] = None # texts input used by ChunkSplit Node + language: str = None # language configuration used by ChunkSplit Node + split_type: str = None # split type used by ChunkSplit Node + example_prompt: str = None # need by graph information extract + schema: str = None # Schema information requeired by SchemaNode + graph_name: str = None + + def reset(self, _: CStatus) -> None: + self.texts = None + self.language = None + self.split_type = None + self.example_prompt = None + self.schema = None + self.graph_name = None + + +class WkFlowState(GParam): + schema: Optional[str] = None # schema message + simple_schema: Optional[str] = None + chunks: Optional[List[str]] = None + edges: Optional[List[Any]] = None + vertices: Optional[List[Any]] = None + triples: Optional[List[Any]] = None + call_count: Optional[int] = None + + keywords: Optional[List[str]] = None + vector_result = None + graph_result = None + keywords_embeddings = None + + def setup(self): + self.schema = None + self.simple_schema = None + self.chunks = None + self.edges = None + self.vertices = None + self.triples = None + self.call_count = None + + self.keywords = None + self.vector_result = None + self.graph_result = None + self.keywords_embeddings = None + + return CStatus() + + def to_json(self): + """ + Automatically returns a JSON-formatted dictionary of all non-None instance members, + eliminating the need to manually maintain the member list. + + Returns: + dict: A dictionary containing non-None instance members and their serialized values. + """ + # Only export instance attributes (excluding methods and class attributes) whose values are not None + return { + k: v + for k, v in self.__dict__.items() + if not k.startswith("_") and v is not None + } diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 9fef06d2b..f61b5f843 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -22,6 +22,7 @@ from typing import Dict, Any, Union, Optional import gradio as gr +from hugegraph_llm.flows.scheduler import SchedulerSingleton from .embedding_utils import get_filename_prefix, get_index_folder_name from .hugegraph_utils import get_hg_client, clean_hg_data @@ -35,11 +36,17 @@ def get_graph_index_info(): - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) graph_summary_info = builder.fetch_graph_data().run() - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) - filename_prefix = get_filename_prefix(llm_settings.embedding_type, getattr(builder.embedding, "model_name", None)) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(builder.embedding, "model_name", None) + ) vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) graph_summary_info["vid_index"] = { "embed_dim": vector_index.index.d, @@ -50,15 +57,20 @@ def get_graph_index_info(): def clean_all_graph_index(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix(llm_settings.embedding_type, - getattr(Embeddings().get_embedding(), "model_name", None)) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, + getattr(Embeddings().get_embedding(), "model_name", None), + ) VectorIndex.clean( - str(os.path.join(resource_path, folder_name, "graph_vids")), - filename_prefix) + str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix + ) VectorIndex.clean( str(os.path.join(resource_path, folder_name, "gremlin_examples")), - filename_prefix) + filename_prefix, + ) log.warning("Clear graph index and text2gql index successfully!") gr.Info("Clear graph index and text2gql index successfully!") @@ -71,7 +83,7 @@ def clean_all_graph_data(): def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: schema = schema.strip() - if schema.startswith('{'): + if schema.startswith("{"): try: schema = json.loads(schema) builder.import_schema(from_user_defined=schema) @@ -84,16 +96,20 @@ def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: return None -def extract_graph(input_file, input_text, schema, example_prompt) -> str: +def extract_graph_origin(input_file, input_text, schema, example_prompt) -> str: texts = read_documents(input_file, input_text) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) if not schema: return "ERROR: please input with correct schema/format." error_message = parse_schema(schema, builder) if error_message: return error_message - builder.chunk_split(texts, "document", "zh").extract_info(example_prompt, "property_graph") + builder.chunk_split(texts, "document", "zh").extract_info( + example_prompt, "property_graph" + ) try: context = builder.run() @@ -103,19 +119,40 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: { "vertices": context["vertices"], "edges": context["edges"], - "warning": "The schema may not match the Doc" + "warning": "The schema may not match the Doc", }, ensure_ascii=False, - indent=2 + indent=2, ) - return json.dumps({"vertices": context["vertices"], "edges": context["edges"]}, ensure_ascii=False, indent=2) + return json.dumps( + {"vertices": context["vertices"], "edges": context["edges"]}, + ensure_ascii=False, + indent=2, + ) + except Exception as e: # pylint: disable=broad-exception-caught + log.error(e) + raise gr.Error(str(e)) + + +def extract_graph(input_file, input_text, schema, example_prompt) -> str: + texts = read_documents(input_file, input_text) + scheduler = SchedulerSingleton.get_instance() + if not schema: + return "ERROR: please input with correct schema/format." + + try: + return scheduler.schedule_flow( + "graph_extract", schema, texts, example_prompt, "property_graph" + ) except Exception as e: # pylint: disable=broad-exception-caught log.error(e) raise gr.Error(str(e)) def update_vid_embedding(): - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) builder.fetch_graph_data().build_vertex_id_semantic_index() log.debug("Operators: %s", builder.operators) try: @@ -132,7 +169,9 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: data_json = json.loads(data.strip()) log.debug("Import graph data: %s", data) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) if schema: error_message = parse_schema(schema, builder) if error_message: @@ -154,7 +193,7 @@ def build_schema(input_text, query_example, few_shot): context = { "raw_texts": [input_text] if input_text else [], "query_examples": [], - "few_shot_schema": {} + "few_shot_schema": {}, } if few_shot: @@ -170,7 +209,7 @@ def build_schema(input_text, query_example, few_shot): context["query_examples"] = [ { "description": ex.get("description", ""), - "gremlin": ex.get("gremlin", "") + "gremlin": ex.get("gremlin", ""), } for ex in parsed_examples if isinstance(ex, dict) and "description" in ex and "gremlin" in ex @@ -178,7 +217,9 @@ def build_schema(input_text, query_example, few_shot): except json.JSONDecodeError as e: raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) try: schema = builder.build_schema().run(context) except Exception as e: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 62bcdd9cb..138b0d359 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -23,11 +23,12 @@ from hugegraph_llm.config import resource_path, huge_settings, llm_settings from hugegraph_llm.indices.vector_index import VectorIndex -from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.kg_construction_task import KgBuilder -from hugegraph_llm.utils.embedding_utils import get_filename_prefix, get_index_folder_name -from hugegraph_llm.utils.hugegraph_utils import get_hg_client +from hugegraph_llm.models.embeddings.init_embedding import model_map +from hugegraph_llm.flows.scheduler import SchedulerSingleton +from hugegraph_llm.utils.embedding_utils import ( + get_filename_prefix, + get_index_folder_name, +) def read_documents(input_file, input_text): @@ -49,7 +50,9 @@ def read_documents(input_file, input_text): texts.append(text) elif full_path.endswith(".pdf"): # TODO: support PDF file - raise gr.Error("PDF will be supported later! Try to upload text/docx now") + raise gr.Error( + "PDF will be supported later! Try to upload text/docx now" + ) else: raise gr.Error("Please input txt or docx file.") else: @@ -59,33 +62,44 @@ def read_documents(input_file, input_text): # pylint: disable=C0301 def get_vector_index_info(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix(llm_settings.embedding_type, - getattr(Embeddings().get_embedding(), "model_name", None)) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) + ) chunk_vector_index = VectorIndex.from_index_file( str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix, - record_miss=False + record_miss=False, ) graph_vid_vector_index = VectorIndex.from_index_file( - str(os.path.join(resource_path, folder_name, "graph_vids")), - filename_prefix + str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix + ) + return json.dumps( + { + "embed_dim": chunk_vector_index.index.d, + "vector_info": { + "chunk_vector_num": chunk_vector_index.index.ntotal, + "graph_vid_vector_num": graph_vid_vector_index.index.ntotal, + "graph_properties_vector_num": len(chunk_vector_index.properties), + }, + }, + ensure_ascii=False, + indent=2, ) - return json.dumps({ - "embed_dim": chunk_vector_index.index.d, - "vector_info": { - "chunk_vector_num": chunk_vector_index.index.ntotal, - "graph_vid_vector_num": graph_vid_vector_index.index.ntotal, - "graph_properties_vector_num": len(chunk_vector_index.properties) - } - }, ensure_ascii=False, indent=2) def clean_vector_index(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - filename_prefix = get_filename_prefix(llm_settings.embedding_type, - getattr(Embeddings().get_embedding(), "model_name", None)) - VectorIndex.clean(str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) + ) + VectorIndex.clean( + str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix + ) gr.Info("Clean vector index successfully!") @@ -93,6 +107,5 @@ def build_vector_index(input_file, input_text): if input_file and input_text: raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - context = builder.chunk_split(texts, "paragraph", "zh").build_vector_index().run() - return json.dumps(context, ensure_ascii=False, indent=2) + scheduler = SchedulerSingleton.get_instance() + return scheduler.schedule_flow("build_vector_index", texts) From b90925ad6d01183657fd91b54f8ed56590475da2 Mon Sep 17 00:00:00 2001 From: Linyu <94553312+weijinglin@users.noreply.github.com> Date: Thu, 25 Sep 2025 15:09:38 +0800 Subject: [PATCH 57/71] refactor: refactor hugegraph-ai to integrate with CGraph & port some usecases in web demo (#49) --- .../spec/hugegraph-llm/fixed_flow/design.md | 643 ++++++++++++++++++ .../hugegraph-llm/fixed_flow/requirements.md | 24 + .../spec/hugegraph-llm/fixed_flow/tasks.md | 36 + .../demo/rag_demo/vector_graph_block.py | 21 +- .../src/hugegraph_llm/flows/build_schema.py | 71 ++ .../hugegraph_llm/flows/build_vector_index.py | 4 +- .../flows/get_graph_index_info.py | 68 ++ .../src/hugegraph_llm/flows/graph_extract.py | 58 +- .../hugegraph_llm/flows/import_graph_data.py | 65 ++ .../hugegraph_llm/flows/prompt_generate.py | 63 ++ .../src/hugegraph_llm/flows/scheduler.py | 31 +- .../flows/update_vid_embeddings.py | 47 ++ .../src/hugegraph_llm/flows/utils.py | 34 + .../src/hugegraph_llm/nodes/base_node.py | 71 ++ .../nodes/document_node/chunk_split.py | 43 ++ .../hugegraph_node/commit_to_hugegraph.py | 35 + .../nodes/hugegraph_node/fetch_graph_data.py | 33 + .../nodes/hugegraph_node/schema.py | 74 ++ .../nodes/index_node/build_semantic_index.py | 34 + .../nodes/index_node/build_vector_index.py | 34 + .../nodes/llm_node/extract_info.py | 52 ++ .../nodes/llm_node/prompt_generate.py | 59 ++ .../nodes/llm_node/schema_build.py | 91 +++ hugegraph-llm/src/hugegraph_llm/nodes/util.py | 27 + .../operators/common_op/check_schema.py | 160 ----- .../operators/document_op/chunk_split.py | 58 -- .../hugegraph_op/commit_to_hugegraph.py | 127 +++- .../operators/hugegraph_op/schema_manager.py | 75 -- .../operators/index_op/build_vector_index.py | 48 -- .../operators/llm_op/info_extract.py | 146 ---- .../llm_op/property_graph_extract.py | 127 +--- .../src/hugegraph_llm/state/ai_state.py | 28 + .../hugegraph_llm/utils/graph_index_utils.py | 40 ++ 33 files changed, 1811 insertions(+), 716 deletions(-) create mode 100644 .vibedev/spec/hugegraph-llm/fixed_flow/design.md create mode 100644 .vibedev/spec/hugegraph-llm/fixed_flow/requirements.md create mode 100644 .vibedev/spec/hugegraph-llm/fixed_flow/tasks.md create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/build_schema.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/utils.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/base_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/document_node/chunk_split.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/commit_to_hugegraph.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/fetch_graph_data.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_semantic_index.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/extract_info.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/util.py diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/design.md b/.vibedev/spec/hugegraph-llm/fixed_flow/design.md new file mode 100644 index 000000000..c5777236d --- /dev/null +++ b/.vibedev/spec/hugegraph-llm/fixed_flow/design.md @@ -0,0 +1,643 @@ +# Hugegraph-ai 固定工作流执行引擎设计文档 + +## 概述 + +Hugegraph固定工作流执行引擎是用来执行固定工作流的工作流执行引擎，每个工作流对应到实际Web Demo的一个具体用例，包括向量索引的构建，图索引的构建等等。该引擎基于PyCGraph框架构建，提供了高性能、可复用的流水线调度能力。 + +### 设计目标 + +- **性能优异**：通过流水线复用机制保证固定工作流的执行性能 +- **高可靠性**：确保数据一致性和故障恢复能力，提供完善的错误处理机制 +- **易于扩展**：能够简单轻松地新增固定工作流，支持动态调度 +- **资源优化**：通过流水线池化管理，减少重复构图开销 + +### 技术栈 + +- **PyCGraph**：基于C++的高性能图计算框架，提供GPipeline和GPipelineManager +- **Python**：主要开发语言，提供业务逻辑和接口层 +- **Threading**：支持并发调度和线程安全 + +### 模块分层 +```text +hugegraph-llm/ +└── src/ + └── hugegraph_llm/ + ├── api/ # FastAPI 接口层，提供 rag_api、admin_api 等服务 + ├── config/ # 配置管理，包含各类配置与生成工具 + ├── demo/ # Gradio Web Demo 及相关交互应用 + ├── document/ # 文档处理与分块等工具 + ├── enums/ # 枚举类型定义 + ├── flows/ # 工作流调度与核心流程（如向量/图索引构建、数据导入等） + │ ├── __init__.py + │ ├── common.py # BaseFlow抽象基类 + │ ├── scheduler.py # 调度器核心实现 + │ ├── build_vector_index.py # 向量索引构建工作流 + │ ├── graph_extract.py # 图抽取工作流 + │ ├── import_graph_data.py # 图数据导入工作流 + │ ├── update_vid_embeddings.py # 向量更新工作流 + │ ├── get_graph_index_info.py # 图索引信息获取工作流 + │ ├── build_schema.py # 模式构建工作流 + │ └── prompt_generate.py # 提示词生成工作流 + ├── indices/ # 各类索引实现（向量、图、关键词等） + ├── middleware/ # 中间件与请求处理 + ├── models/ # LLM、Embedding、Reranker 等模型相关 + ├── nodes/ # Node调度层，负责Operator生命周期和上下文管理 + │ ├── base_node.py + │ ├── document_node/ + │ ├── hugegraph_node/ + │ ├── index_node/ + │ ├── llm_node/ + │ └── util.py + ├── operators/ # 主要算子与任务（如 KG 构建、GraphRAG、Text2Gremlin 等） + ├── resources/ # 资源文件（Prompt、示例、Gremlin 模板等） + ├── state/ # 状态管理 + ├── utils/ # 工具类与通用方法 + └── __init__.py # 包初始化 +``` + +## 架构设计 + +### 整体架构 + +> 新架构在Flow与Operator之间引入Node层，Node负责Operator的生命周期管理、上下文绑定、参数区解耦和并发安全，所有Flow均通过Node组装，Operator只关注业务实现。 + +#### 架构图 + +```mermaid +graph TB + subgraph UserLayer["用户层"] + User["用户请求"] + end + + subgraph SchedulerLayer["调度层"] + Scheduler["Scheduler
调度器"] + Singleton["SchedulerSingleton
单例管理器"] + end + + subgraph FlowLayer["工作流层"] + Pool["pipeline_pool
流水线池"] + BVI["BuildVectorIndexFlow
向量索引构建"] + GE["GraphExtractFlow
图抽取工作流"] + end + + subgraph PyCGraphLayer["PyCGraph层"] + Manager1["GPipelineManager
向量索引管理器"] + Manager2["GPipelineManager
图抽取管理器"] + Pipeline1["GPipeline
向量索引流水线"] + Pipeline2["GPipeline
图抽取流水线"] + end + + subgraph OperatorLayer["算子层"] + ChunkSplit["ChunkSplitNode
文档分块"] + BuildVector["BuildVectorIndexNode
向量索引构建"] + SchemaNode["SchemaNode
模式管理"] + InfoExtract["ExtractNode
信息抽取"] + PropGraph["Commit2GraphNode
图数据导入"] + FetchNode["FetchGraphDataNode
图数据拉取"] + SemanticIndex["BuildSemanticIndexNode
语义索引构建"] + end + + subgraph StateLayer["状态层"] + WkInput["wkflow_input
工作流输入"] + WkState["wkflow_state
工作流状态"] + end + + User --> Scheduler + Scheduler --> Singleton + Scheduler --> Pool + Pool --> BVI + Pool --> GE + BVI --> Manager1 + GE --> Manager2 + Manager1 --> Pipeline1 + Manager2 --> Pipeline2 + Pipeline1 --> ChunkSplit + Pipeline1 --> BuildVector + Pipeline2 --> SchemaNode + Pipeline2 --> ChunkSplit + Pipeline2 --> InfoExtract + Pipeline2 --> PropGraph + Pipeline1 --> WkInput + Pipeline1 --> WkState + Pipeline2 --> WkInput + Pipeline2 --> WkState + + style Scheduler fill:#e1f5fe + style Pool fill:#f3e5f5 + style Manager1 fill:#fff3e0 + style Manager2 fill:#fff3e0 + style Pipeline1 fill:#e8f5e8 + style Pipeline2 fill:#e8f5e8 +``` + +#### 调度流程图 + +```mermaid +flowchart TD + Start([开始]) --> CheckFlow{检查工作流
是否支持} + CheckFlow -->|否| Error1[抛出ValueError] + CheckFlow -->|是| FetchPipeline[从Manager获取
可复用Pipeline] + + FetchPipeline --> IsNull{Pipeline
是否为null} + + IsNull -->|是| BuildNew[构建新Pipeline] + BuildNew --> InitPipeline[初始化Pipeline] + InitPipeline --> InitCheck{初始化
是否成功} + InitCheck -->|否| Error2[记录错误并中止] + InitCheck -->|是| RunPipeline[执行Pipeline] + RunPipeline --> RunCheck{执行
是否成功} + RunCheck -->|否| Error3[记录错误并中止] + RunCheck -->|是| PostDeal[后处理结果] + PostDeal --> AddToPool[添加到复用池] + AddToPool --> Return[返回结果] + + IsNull -->|否| PrepareInput[准备输入数据] + PrepareInput --> RunReused[执行复用Pipeline] + RunReused --> ReusedCheck{执行
是否成功} + ReusedCheck -->|否| Error4[抛出RuntimeError] + ReusedCheck -->|是| PostDealReused[后处理结果] + PostDealReused --> ReleasePipeline[释放Pipeline] + ReleasePipeline --> Return + + Error1 --> End([结束]) + Error2 --> End + Error3 --> End + Error4 --> End + Return --> End + + style Start fill:#4caf50 + style End fill:#f44336 + style CheckFlow fill:#ff9800 + style IsNull fill:#ff9800 + style InitCheck fill:#ff9800 + style RunCheck fill:#ff9800 + style ReusedCheck fill:#ff9800 +``` + +### 核心组件 + +#### 1. Scheduler（调度器） +- **职责**：调度中心，维护 `pipeline_pool`，提供统一的工作流调度接口 +- **特性**： + - 支持多种工作流类型（build_vector_index、graph_extract、import_graph_data、update_vid_embeddings、get_graph_index_info、build_schema、prompt_generate等） + - 流水线池化管理，支持复用 + - 线程安全的单例模式 + - 可配置的最大流水线数量 + +#### 2. GPipelineManager（流水线管理器） +- **来源**：PyCGraph框架提供 +- **职责**：负责流水线对象 `GPipeline` 的获取、添加、释放与复用 +- **特性**： + - 自动管理流水线生命周期 + - 支持流水线复用和资源回收 + - 提供fetch/add/release操作接口 + +#### 3. BaseFlow（工作流基类） +- **职责**：工作流构建与前后处理抽象 +- **接口**： + - `prepare()`: 预处理接口，准备输入数据 + - `build_flow()`: 组装Node并注册依赖关系 + - `post_deal()`: 后处理接口，处理执行结果 +- **实现**： + - `BuildVectorIndexFlow`: 向量索引构建工作流 + - `GraphExtractFlow`: 图抽取工作流 + - `ImportGraphDataFlow`: 图数据导入工作流 + - `UpdateVidEmbeddingsFlows`: 向量更新工作流 + - `GetGraphIndexInfoFlow`: 图索引信息获取工作流 + - `BuildSchemaFlow`: 模式构建工作流 + - `PromptGenerateFlow`: 提示词生成工作流 + +#### 4. Node（节点调度器） +- **职责**：作为Operator的生命周期管理者，负责参数区绑定、上下文初始化、并发安全、异常处理等。 +- **特性**： + - 统一生命周期接口（init、node_init、run、operator_schedule） + - 通过参数区（wkflow_input/wkflow_state）与Flow/Operator解耦 + - Operator只需实现run(data_json)方法，Node负责调度和结果写回 + - 典型Node如：ChunkSplitNode、BuildVectorIndexNode、SchemaNode、ExtractNode、Commit2GraphNode、FetchGraphDataNode、BuildSemanticIndexNode、SchemaBuildNode、PromptGenerateNode等 + +#### 5. Operator（算子） +- **职责**：实现具体的业务原子操作 +- **特性**： + - 只需关注自身业务逻辑实现 + - 由Node统一调度 + +#### 6. GPipeline（流水线实例） +- **来源**：PyCGraph框架提供 +- **职责**：具体流水线实例，包含参数区与节点DAG拓扑 +- **参数区**： + - `wkflow_input`: 流水线运行输入 + - `wkflow_state`: 流水线运行状态与中间结果 + +### 核心数据结构 + +```python +# Scheduler核心数据结构 +Scheduler.pipeline_pool: Dict[str, Any] = { + "build_vector_index": { + "manager": GPipelineManager(), + "flow": BuildVectorIndexFlow(), + }, + "graph_extract": { + "manager": GPipelineManager(), + "flow": GraphExtractFlow(), + } +} +``` + +### 调度流程 + +#### schedule_flow方法执行流程 + +1. **工作流验证**：校验 `flow` 是否受支持，查表获取对应的 `manager` 与 `flow` 实例 + +2. **流水线获取**：从 `manager.fetch()` 获取可复用的 `GPipeline` + +3. **新流水线处理**（当fetch()返回None时）： + - 调用 `flow.build_flow(*args, **kwargs)` 构建新流水线 + - 调用 `pipeline.init()` 完成初始化，失败则记录错误并中止 + - 调用 `pipeline.run()` 执行，失败则中止 + - 调用 `flow.post_deal(pipeline)` 生成输出 + - 调用 `manager.add(pipeline)` 将流水线加入可复用池 + +4. **复用流水线处理**（当fetch()返回现有流水线时）： + - 从 `pipeline.getGParamWithNoEmpty("wkflow_input")` 获取输入对象 + - 调用 `flow.prepare(prepared_input, *args, **kwargs)` 进行参数刷新 + - 调用 `pipeline.run()` 执行，失败则中止 + - 调用 `flow.post_deal(pipeline)` 生成输出 + - 调用 `manager.release(pipeline)` 归还流水线 + +### 并发与复用策略 + +#### 线程安全 +- `SchedulerSingleton` 使用双重检查锁保证全局单例 +- 线程安全获取 `Scheduler` 实例 + +#### 资源管理 +- 每种 `flow` 拥有独立的 `GPipelineManager` +- 最大并发量由 `Scheduler.max_pipeline` 与底层 `GPipelineManager` 策略共同约束 +- 通过 `fetch/add/release` 机制减少重复构图的开销 + +#### 性能优化 +- 流水线复用机制适合高频相同工作流场景 +- 减少重复初始化和构图的时间开销 +- 支持并发执行多个工作流实例 + +### 错误处理与日志 + +#### 错误检测 +- 对 `init/run` 的 `Status.isErr()` 进行检测 +- 统一抛出 `RuntimeError` 并记录详细 `status.getInfo()` +- 提供完整的错误堆栈信息 + +#### 日志记录 +- 使用统一的日志系统记录关键操作 +- 记录流水线执行状态和错误信息 +- 支持不同级别的日志输出 + +#### 结果处理 +- `flow.post_deal` 负责将 `wkflow_state` 转换为对外可消费结果（如JSON） +- 提供标准化的输出格式 +- 支持错误信息的友好展示 + +### 扩展指引 + +#### 新增Node/Operator/Flow步骤 +1. 实现Operator业务逻辑（如ChunkSplit/BuildVectorIndex/InfoExtract等） +2. 实现对应Node（继承BaseNode，负责参数区绑定和调度Operator） +3. 在Flow中组装Node，注册依赖关系 +4. 在Scheduler注册新的Flow + +#### 输入输出约定 +- 统一使用 `wkflow_input` 作为输入载体 +- 统一使用 `wkflow_state` 作为状态与结果容器 +- 确保可复用流水线在不同请求间可被快速重置 + +#### 最佳实践 +- 保持Flow类的无状态设计 +- 合理使用流水线复用机制 +- 提供完善的错误处理和日志记录 +- 遵循统一的接口规范 + +## Flow对象设计 + +### BaseFlow抽象基类 + +```python +class BaseFlow(ABC): + """ + Base class for flows, defines three interface methods: prepare, build_flow, and post_deal. + """ + + @abstractmethod + def prepare(self, prepared_input: WkFlowInput, *args, **kwargs): + """ + Pre-processing interface. + """ + pass + + @abstractmethod + def build_flow(self, *args, **kwargs): + """ + Interface for building the flow. + """ + pass + + @abstractmethod + def post_deal(self, *args, **kwargs): + """ + Post-processing interface. + """ + pass +``` + +### 接口说明 + +每个Flow对象都需要实现三个核心接口： + +- **prepare**: 用来准备整个workflow的输入数据，设置工作流参数 +- **build_flow**: 用来构建整个workflow的流水线，注册节点和依赖关系 +- **post_deal**: 用来处理workflow的执行结果，转换为对外输出格式 + +### 具体实现示例 + +#### BuildVectorIndexFlow（向量索引构建工作流） + +```python +class BuildVectorIndexFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, texts): + prepared_input.texts = texts + prepared_input.language = "zh" + prepared_input.split_type = "paragraph" + return + + def build_flow(self, texts): + pipeline = GPipeline() + # prepare for workflow input + prepared_input = WkFlowInput() + self.prepare(prepared_input, texts) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + chunk_split_node = ChunkSplitNode() + build_vector_node = BuildVectorIndexNode() + pipeline.registerGElement(chunk_split_node, set(), "chunk_split") + pipeline.registerGElement(build_vector_node, {chunk_split_node}, "build_vector") + + return pipeline + + def post_deal(self, pipeline=None): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + return json.dumps(res, ensure_ascii=False, indent=2) +``` + +#### GraphExtractFlow（图抽取工作流） + +```python +class GraphExtractFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type): + prepared_input.texts = texts + prepared_input.language = "zh" + prepared_input.split_type = "document" + prepared_input.example_prompt = example_prompt + prepared_input.schema = schema + prepare_schema(prepared_input, schema) + return + + def build_flow(self, schema, texts, example_prompt, extract_type): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, schema, texts, example_prompt, extract_type) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + schema_node = SchemaNode() + + chunk_split_node = ChunkSplitNode() + graph_extract_node = ExtractNode() + + pipeline.registerGElement(schema_node, set(), "schema_node") + pipeline.registerGElement(chunk_split_node, set(), "chunk_split") + pipeline.registerGElement( + graph_extract_node, {schema_node, chunk_split_node}, "graph_extract" + ) + + return pipeline + + def post_deal(self, pipeline=None): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + vertices = res.get("vertices", []) + edges = res.get("edges", []) + if not vertices and not edges: + log.info("Please check the schema.(The schema may not match the Doc)") + return json.dumps( + { + "vertices": vertices, + "edges": edges, + "warning": "The schema may not match the Doc", + }, + ensure_ascii=False, + indent=2, + ) + return json.dumps( + {"vertices": vertices, "edges": edges}, + ensure_ascii=False, + indent=2, + ) +``` + +## Node对象设计 + +### 节点生命周期 + +节点以 GNode 为抽象基类，统一生命周期与状态返回。方法职责与约定如下： + +#### 初始化阶段 + +- **init()**: + - **责任**：完成节点级初始化工作（如绑定共享上下文、准备参数区），确保节点具备运行所需的最小环境 + - **约定**：仅做轻量初始化，不执行业务逻辑；返回状态用于判断是否可继续 + +- **node_init()**: + - **责任**：解析与校验本次运行所需的输入（通常来自 wk_input），构建运行期依赖（如内部配置、变换器、资源句柄） + - **约定**：输入缺失或不合法时，应返回错误状态并中止后续执行；不产生对外可见的业务结果 + +#### 运行阶段 + +- **run()**: + - **责任**：执行业务主流程（纯计算或 I/O），在完成后将节点产出写入共享状态（wkflow_state/上下文） + - **约定**： + - 进入前应先调用 node_init() 并检查其返回状态 + - 对共享状态的写操作需遵循并发安全约定（如加锁/解锁） + - 出错使用统一状态返回，不抛出未捕获异常到流程编排层 + +### 输入/输出与上下文约定 + +- **输入**：通过编排层预置于参数区（如 wk_input），节点在 node_init() 中读取并校验 +- **输出**：通过共享状态容器（如 wkflow_state/上下文）对外暴露，键/字段命名应稳定可预期，供下游节点消费 + +### 错误处理约定 + +- 统一以状态对象表示成功/失败与信息；错误应尽早返回，避免在 run() 中继续副作用操作 +- 对可预见的校验类错误使用明确的错误信息，便于定位问题与编排层记录 + +### 并发与可重入约定 + +- 共享状态的写入需在临界区内完成；读取视数据一致性要求决定是否加锁 +- 节点应尽量保持无副作用或将副作用范围收敛在可控区域，以支持重试与复用 + +### 可测试性与解耦 + +- 业务纯逻辑应与框架交互解耦，优先封装为可单测的纯函数/内部方法 +- 节点仅负责生命周期编排与上下文读写，具体策略与算法通过内部可替换组件提供 + +### 节点类型 + +#### 文档处理节点 +- **ChunkSplitNode**: 文档分块处理节点 + - 功能：将输入文档按照指定策略进行分块 + - 输入：原始文档文本 + - 输出：分块后的文档片段 + +#### 索引构建节点 +- **BuildVectorIndexNode**: 向量索引构建节点 + - 功能：基于文档分块构建向量索引 + - 输入：文档分块 + - 输出：向量索引数据 + +#### 模式管理节点 +- **SchemaManagerNode**: 图模式管理节点 + - 功能：从HugeGraph获取图模式信息 + - 输入：图名称 + - 输出：图模式定义 + +- **CheckSchemaNode**: 模式校验节点 + - 功能：校验用户定义的图模式 + - 输入：用户定义的JSON模式 + - 输出：校验后的模式定义 + +#### 图抽取节点 +- **InfoExtractNode**: 信息抽取节点 + - 功能：从文档中抽取三元组信息 + - 输入：文档分块和模式定义 + - 输出：抽取的三元组数据 + +- **PropertyGraphExtractNode**: 属性图抽取节点 + - 功能：从文档中抽取属性图结构 + - 输入：文档分块和模式定义 + - 输出：抽取的顶点和边数据 + +#### 模式构建节点 +- **SchemaBuildNode**: 模式构建节点 + - 功能：基于文档和查询示例构建图模式 + - 输入：文档文本、查询示例、少样本模式 + - 输出：构建的图模式定义 + +#### 提示词生成节点 +- **PromptGenerateNode**: 提示词生成节点 + - 功能：基于源文本、场景和示例名称生成提示词 + - 输入：源文本、场景、示例名称 + - 输出：生成的提示词 + + +## 测试策略 + +### 测试目标 + +目前的测试策略主要目标是保证移植之后的workflow和移植之前的workflow执行结果、程序行为一致。 + +### 测试范围 + +#### 1. 功能测试 +- **工作流执行结果一致性**：确保新架构下的工作流执行结果与原有实现完全一致 +- **输入输出格式验证**：验证输入参数处理和输出格式转换的正确性 +- **错误处理测试**：确保错误场景下的行为与预期一致 + +#### 2. 性能测试 +- **流水线复用效果**：验证流水线复用机制的性能提升效果 +- **并发执行测试**：测试多工作流并发执行的稳定性和性能 +- **资源使用测试**：监控内存和CPU使用情况，确保资源使用合理 + +#### 3. 稳定性测试 +- **长时间运行测试**：验证系统在长时间运行下的稳定性 +- **异常恢复测试**：测试系统在异常情况下的恢复能力 +- **内存泄漏测试**：确保流水线复用不会导致内存泄漏 + +### 测试方法 + +#### 1. 单元测试 +- 对每个Flow类进行单元测试 +- 对每个Node类进行单元测试 +- 对Scheduler调度逻辑进行测试 + +#### 2. 集成测试 +- 端到端工作流测试 +- 多工作流组合测试 +- 与外部系统集成测试 + +#### 3. 性能基准测试 +- 建立性能基准线 +- 对比新旧架构的性能差异 +- 监控关键性能指标 + +### 测试数据 + +#### 1. 标准测试数据集 +- 准备标准化的测试文档 +- 准备标准化的图模式定义 +- 准备标准化的期望输出结果 + +#### 2. 边界测试数据 +- 空输入测试 +- 大文件测试 +- 特殊字符测试 +- 异常格式测试 + +### 测试环境 + +#### 1. 开发环境测试 +- 本地开发环境的功能验证 +- 快速迭代测试 + +#### 2. 测试环境验证 +- 模拟生产环境的完整测试 +- 性能压力测试 + +#### 3. 生产环境验证 +- 灰度发布验证 +- 生产环境监控 + +### 测试自动化 + +#### 1. CI/CD集成 +- 自动化测试流程集成 +- 代码提交触发测试 +- 测试结果自动报告 + +#### 2. 回归测试 +- 定期执行回归测试 +- 确保新功能不影响现有功能 +- 性能回归检测 + +### 测试指标 + +#### 1. 功能指标 +- 测试覆盖率 > 90% +- 功能正确性 100% +- 错误处理覆盖率 > 95% + +#### 2. 性能指标 +- 响应时间提升 > 20% +- 吞吐量提升 > 30% +- 资源使用优化 > 15% + +#### 3. 稳定性指标 +- 系统可用性 > 99.9% +- 平均故障恢复时间 < 5分钟 +- 内存泄漏率 = 0% diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/requirements.md b/.vibedev/spec/hugegraph-llm/fixed_flow/requirements.md new file mode 100644 index 000000000..095027369 --- /dev/null +++ b/.vibedev/spec/hugegraph-llm/fixed_flow/requirements.md @@ -0,0 +1,24 @@ +## 需求列表 + +### 核心框架设计 + +**核心**：Scheduler类中的schedule_flow设计与实现 + +**验收标准**： +1.1. 核心框架尽可能复用资源，避免资源的重复分配和释放 +1.2. 应该保证正常的请求处理指标要求 +1.3. 应该能够配置框架整体使用的资源上限 + +### 固定工作流移植 + +**核心**：移植Web Demo中的所有用例 +2.1. 保证使用核心框架移植后的工作流的程序行为和移植之前保持一致即可 + +**已完成的工作流类型**： +- build_vector_index: 向量索引构建工作流 +- graph_extract: 图抽取工作流 +- import_graph_data: 图数据导入工作流 +- update_vid_embeddings: 向量更新工作流 +- get_graph_index_info: 图索引信息获取工作流 +- build_schema: 模式构建工作流 +- prompt_generate: 提示词生成工作流 diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/tasks.md b/.vibedev/spec/hugegraph-llm/fixed_flow/tasks.md new file mode 100644 index 000000000..a84aee2ff --- /dev/null +++ b/.vibedev/spec/hugegraph-llm/fixed_flow/tasks.md @@ -0,0 +1,36 @@ +# HugeGraph-ai 固定工作流框架设计和用例移植 + +本文档将 HugeGraph 固定工作流框架设计和用例移植转换为一系列可执行的编码任务。 + +## 1. schedule_flow设计与实现 + +- [x] **1.1 构建Scheduler框架1.0** + - 需要能够复用已经创建过的Pipeline(Pipeline Pooling) + - 使用CGraph(Graph-based engine)作为底层执行引擎 + - 不同Node之间松耦合 + +- [ ] **1.2 优化Scheduler框架资源配置** + - 支持用户配置底层线程池参数 + - 现有的workflow可能会根据输入有细小的变化，导致相同的用例得到不同的workflow，怎么解决这个问题呢？ + - Node/Operator解耦，Node负责生命周期和上下文，Operator只关注业务逻辑 + - Flow只负责组装Node，所有业务逻辑下沉到Node/Operator + - Scheduler支持多类型Flow注册，注册方式更灵活 + +- [ ] **1.3 优化Scheduler框架资源使用** + - 根据负载控制每个PipelineManager管理的Pipeline数量，实现动态扩缩容 + - Node层支持参数区自动绑定和并发安全 + - Operator只需实现run(data_json)方法，Node负责调度和结果写回 + +## 2. 固定工作流用例移植 + +- [x] **2.1 build_vector_index workflow移植** +- [x] **2.2 graph_extract workflow移植** +- [x] **2.3 import_graph_data workflow移植** + - 基于Node/Operator机制实现import_graph_data工作流 +- [x] **2.4 update_vid_embeddings workflow移植** + - 基于Node/Operator机制实现update_vid_embeddings工作流 +- [x] **2.5 get_graph_index_info workflow移植** +- [x] **2.6 build_schema workflow移植** + - 基于Node/Operator机制实现build_schema工作流 +- [x] **2.7 prompt_generate workflow移植** + - 基于Node/Operator机制实现prompt_generate工作流 diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 9897f420f..4aa476942 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -26,8 +26,7 @@ from hugegraph_llm.config import huge_settings from hugegraph_llm.config import prompt from hugegraph_llm.config import resource_path -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.llm_op.prompt_generate import PromptGenerate +from hugegraph_llm.flows.scheduler import SchedulerSingleton from hugegraph_llm.utils.graph_index_utils import ( get_graph_index_info, clean_all_graph_index, @@ -61,7 +60,7 @@ def store_prompt(doc, schema, example_prompt): def generate_prompt_for_ui(source_text, scenario, example_name): """ - Handles the UI logic for generating a new prompt. It calls the PromptGenerate operator. + Handles the UI logic for generating a new prompt using the new workflow architecture. """ if not all([source_text, scenario, example_name]): gr.Warning( @@ -69,19 +68,13 @@ def generate_prompt_for_ui(source_text, scenario, example_name): ) return gr.update() try: - prompt_generator = PromptGenerate(llm=LLMs().get_chat_llm()) - context = { - "source_text": source_text, - "scenario": scenario, - "example_name": example_name, - } - result_context = prompt_generator.run(context) - # Presents the result of generating prompt - generated_prompt = result_context.get( - "generated_extract_prompt", "Generation failed. Please check the logs." + # using new architecture + scheduler = SchedulerSingleton.get_instance() + result = scheduler.schedule_flow( + "prompt_generate", source_text, scenario, example_name ) gr.Info("Prompt generated successfully!") - return generated_prompt + return result except Exception as e: log.error("Error generating Prompt: %s", e, exc_info=True) raise gr.Error(f"Error generating Prompt: {e}") from e diff --git a/hugegraph-llm/src/hugegraph_llm/flows/build_schema.py b/hugegraph-llm/src/hugegraph_llm/flows/build_schema.py new file mode 100644 index 000000000..6bbcb8512 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/build_schema.py @@ -0,0 +1,71 @@ +# 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.common import BaseFlow +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.nodes.llm_node.schema_build import SchemaBuildNode +from hugegraph_llm.utils.log import log + +import json +from PyCGraph import GPipeline + + +class BuildSchemaFlow(BaseFlow): + def __init__(self): + pass + + def prepare( + self, + prepared_input: WkFlowInput, + texts=None, + query_examples=None, + few_shot_schema=None, + ): + prepared_input.texts = texts + # Optional fields packed into wk_input for SchemaBuildNode + # Keep raw values; node will parse if strings + prepared_input.query_examples = query_examples + prepared_input.few_shot_schema = few_shot_schema + return + + def build_flow(self, texts=None, query_examples=None, few_shot_schema=None): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare( + prepared_input, + texts=texts, + query_examples=query_examples, + few_shot_schema=few_shot_schema, + ) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + schema_build_node = SchemaBuildNode() + pipeline.registerGElement(schema_build_node, set(), "schema_build") + + return pipeline + + def post_deal(self, pipeline=None): + state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + if "schema" not in state_json: + return "" + res = state_json["schema"] + try: + formatted_schema = json.dumps(res, ensure_ascii=False, indent=2) + return formatted_schema + except (TypeError, ValueError) as e: + log.error("Failed to format schema: %s", e) + return str(res) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py index f1ee8c1c4..9a07b5dba 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py @@ -14,13 +14,13 @@ # limitations under the License. from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.document_node.chunk_split import ChunkSplitNode +from hugegraph_llm.nodes.index_node.build_vector_index import BuildVectorIndexNode from hugegraph_llm.state.ai_state import WkFlowInput import json from PyCGraph import GPipeline -from hugegraph_llm.operators.document_op.chunk_split import ChunkSplitNode -from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndexNode from hugegraph_llm.state.ai_state import WkFlowState diff --git a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py new file mode 100644 index 000000000..fa10d0199 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py @@ -0,0 +1,68 @@ +# 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 os + +from hugegraph_llm.config import huge_settings, llm_settings, resource_path +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.indices.vector_index import VectorIndex +from hugegraph_llm.models.embeddings.init_embedding import model_map +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.nodes.hugegraph_node.fetch_graph_data import FetchGraphDataNode +from PyCGraph import GPipeline +from hugegraph_llm.utils.embedding_utils import ( + get_filename_prefix, + get_index_folder_name, +) + + +class GetGraphIndexInfoFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, *args, **kwargs): + return + + def build_flow(self, *args, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, *args, **kwargs) + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + fetch_node = FetchGraphDataNode() + pipeline.registerGElement(fetch_node, set(), "fetch_node") + return pipeline + + def post_deal(self, pipeline=None): + graph_summary_info = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) + index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) + filename_prefix = get_filename_prefix( + llm_settings.embedding_type, + model_map.get(llm_settings.embedding_type, None), + ) + try: + vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) + except FileNotFoundError: + return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) + graph_summary_info["vid_index"] = { + "embed_dim": vector_index.index.d, + "num_vectors": vector_index.index.ntotal, + "num_vids": len(vector_index.properties), + } + return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py index f1a6c5f6f..1b0c98253 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py @@ -16,14 +16,10 @@ import json from PyCGraph import GPipeline from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.document_node.chunk_split import ChunkSplitNode +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.nodes.llm_node.extract_info import ExtractNode from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState -from hugegraph_llm.operators.common_op.check_schema import CheckSchemaNode -from hugegraph_llm.operators.document_op.chunk_split import ChunkSplitNode -from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManagerNode -from hugegraph_llm.operators.llm_op.info_extract import InfoExtractNode -from hugegraph_llm.operators.llm_op.property_graph_extract import ( - PropertyGraphExtractNode, -) from hugegraph_llm.utils.log import log @@ -31,21 +27,6 @@ class GraphExtractFlow(BaseFlow): def __init__(self): pass - def _import_schema( - self, - from_hugegraph=None, - from_extraction=None, - from_user_defined=None, - ): - if from_hugegraph: - return SchemaManagerNode() - elif from_user_defined: - return CheckSchemaNode() - elif from_extraction: - raise NotImplementedError("Not implemented yet") - else: - raise ValueError("No input data / invalid schema type") - def prepare( self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type ): @@ -55,17 +36,7 @@ def prepare( prepared_input.split_type = "document" prepared_input.example_prompt = example_prompt prepared_input.schema = schema - schema = schema.strip() - if schema.startswith("{"): - try: - schema = json.loads(schema) - prepared_input.schema = schema - except json.JSONDecodeError as exc: - log.error("Invalid JSON format in schema. Please check it again.") - raise ValueError("Invalid JSON format in schema.") from exc - else: - log.info("Get schema '%s' from graphdb.", schema) - prepared_input.graph_name = schema + prepared_input.extract_type = extract_type return def build_flow(self, schema, texts, example_prompt, extract_type): @@ -76,27 +47,10 @@ def build_flow(self, schema, texts, example_prompt, extract_type): pipeline.createGParam(prepared_input, "wkflow_input") pipeline.createGParam(WkFlowState(), "wkflow_state") - schema = schema.strip() - schema_node = None - if schema.startswith("{"): - try: - schema = json.loads(schema) - schema_node = self._import_schema(from_user_defined=schema) - except json.JSONDecodeError as exc: - log.error("Invalid JSON format in schema. Please check it again.") - raise ValueError("Invalid JSON format in schema.") from exc - else: - log.info("Get schema '%s' from graphdb.", schema) - schema_node = self._import_schema(from_hugegraph=schema) + schema_node = SchemaNode() chunk_split_node = ChunkSplitNode() - graph_extract_node = None - if extract_type == "triples": - graph_extract_node = InfoExtractNode() - elif extract_type == "property_graph": - graph_extract_node = PropertyGraphExtractNode() - else: - raise ValueError(f"Unsupported extract_type: {extract_type}") + graph_extract_node = ExtractNode() pipeline.registerGElement(schema_node, set(), "schema_node") pipeline.registerGElement(chunk_split_node, set(), "chunk_split") pipeline.registerGElement( diff --git a/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py b/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py new file mode 100644 index 000000000..5581ef107 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py @@ -0,0 +1,65 @@ +# 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 gradio as gr +from PyCGraph import GPipeline +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.hugegraph_node.commit_to_hugegraph import Commit2GraphNode +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.utils.log import log + + +class ImportGraphDataFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, data, schema): + try: + data_json = json.loads(data.strip()) if isinstance(data, str) else data + except json.JSONDecodeError as e: + raise ValueError(f"Invalid JSON for 'data': {e.msg}") from e + log.debug( + "Import graph data (truncated): %s", + (data[:512] + "...") + if isinstance(data, str) and len(data) > 512 + else (data if isinstance(data, str) else ""), + ) + prepared_input.data_json = data_json + prepared_input.schema = schema + return + + def build_flow(self, data, schema): + pipeline = GPipeline() + prepared_input = WkFlowInput() + # prepare input data + self.prepare(prepared_input, data, schema) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + schema_node = SchemaNode() + commit_node = Commit2GraphNode() + pipeline.registerGElement(schema_node, set(), "schema_node") + pipeline.registerGElement(commit_node, {schema_node}, "commit_node") + + return pipeline + + def post_deal(self, pipeline=None): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + gr.Info("Import graph data successfully!") + return json.dumps(res, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py new file mode 100644 index 000000000..aece6bd61 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py @@ -0,0 +1,63 @@ +# 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.common import BaseFlow +from hugegraph_llm.nodes.llm_node.prompt_generate import PromptGenerateNode +from hugegraph_llm.state.ai_state import WkFlowInput + +from PyCGraph import GPipeline + +from hugegraph_llm.state.ai_state import WkFlowState + + +class PromptGenerateFlow(BaseFlow): + def __init__(self): + pass + + def prepare(self, prepared_input: WkFlowInput, source_text, scenario, example_name): + """ + Prepare input data for PromptGenerate workflow + """ + prepared_input.source_text = source_text + prepared_input.scenario = scenario + prepared_input.example_name = example_name + return + + def build_flow(self, source_text, scenario, example_name): + """ + Build the PromptGenerate workflow + """ + pipeline = GPipeline() + # Prepare workflow input + prepared_input = WkFlowInput() + self.prepare(prepared_input, source_text, scenario, example_name) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + # Create PromptGenerate node + prompt_generate_node = PromptGenerateNode() + pipeline.registerGElement(prompt_generate_node, set(), "prompt_generate") + + return pipeline + + def post_deal(self, pipeline=None): + """ + Process the execution result of PromptGenerate workflow + """ + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + return res.get( + "generated_extract_prompt", "Generation failed. Please check the logs." + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index b096310db..559540ce3 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -15,10 +15,15 @@ import threading from typing import Dict, Any -from PyCGraph import GPipelineManager +from PyCGraph import GPipeline, GPipelineManager from hugegraph_llm.flows.build_vector_index import BuildVectorIndexFlow from hugegraph_llm.flows.common import BaseFlow from hugegraph_llm.flows.graph_extract import GraphExtractFlow +from hugegraph_llm.flows.import_graph_data import ImportGraphDataFlow +from hugegraph_llm.flows.update_vid_embeddings import UpdateVidEmbeddingsFlows +from hugegraph_llm.flows.get_graph_index_info import GetGraphIndexInfoFlow +from hugegraph_llm.flows.build_schema import BuildSchemaFlow +from hugegraph_llm.flows.prompt_generate import PromptGenerateFlow from hugegraph_llm.utils.log import log @@ -37,6 +42,26 @@ def __init__(self, max_pipeline: int = 10): "manager": GPipelineManager(), "flow": GraphExtractFlow(), } + self.pipeline_pool["import_graph_data"] = { + "manager": GPipelineManager(), + "flow": ImportGraphDataFlow(), + } + self.pipeline_pool["update_vid_embeddings"] = { + "manager": GPipelineManager(), + "flow": UpdateVidEmbeddingsFlows(), + } + self.pipeline_pool["get_graph_index_info"] = { + "manager": GPipelineManager(), + "flow": GetGraphIndexInfoFlow(), + } + self.pipeline_pool["build_schema"] = { + "manager": GPipelineManager(), + "flow": BuildSchemaFlow(), + } + self.pipeline_pool["prompt_generate"] = { + "manager": GPipelineManager(), + "flow": PromptGenerateFlow(), + } self.max_pipeline = max_pipeline # TODO: Implement Agentic Workflow @@ -46,9 +71,9 @@ def agentic_flow(self): def schedule_flow(self, flow: str, *args, **kwargs): if flow not in self.pipeline_pool: raise ValueError(f"Unsupported workflow {flow}") - manager = self.pipeline_pool[flow]["manager"] + manager: GPipelineManager = self.pipeline_pool[flow]["manager"] flow: BaseFlow = self.pipeline_pool[flow]["flow"] - pipeline = manager.fetch() + pipeline: GPipeline = manager.fetch() if pipeline is None: # call coresponding flow_func to create new workflow pipeline = flow.build_flow(*args, **kwargs) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py b/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py new file mode 100644 index 000000000..b3f0d9923 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py @@ -0,0 +1,47 @@ +# 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 PyCGraph import CStatus, GPipeline +from hugegraph_llm.flows.common import BaseFlow, WkFlowInput +from hugegraph_llm.nodes.hugegraph_node.fetch_graph_data import FetchGraphDataNode +from hugegraph_llm.nodes.index_node.build_semantic_index import BuildSemanticIndexNode +from hugegraph_llm.state.ai_state import WkFlowState + + +class UpdateVidEmbeddingsFlows(BaseFlow): + def prepare(self, prepared_input: WkFlowInput): + return CStatus() + + def build_flow(self): + pipeline = GPipeline() + prepared_input = WkFlowInput() + # prepare input data + self.prepare(prepared_input) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + fetch_node = FetchGraphDataNode() + build_node = BuildSemanticIndexNode() + pipeline.registerGElement(fetch_node, set(), "fetch_node") + pipeline.registerGElement(build_node, {fetch_node}, "build_node") + + return pipeline + + def post_deal(self, pipeline): + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + removed_num = res.get("removed_vid_vector_num", 0) + added_num = res.get("added_vid_vector_num", 0) + return f"Removed {removed_num} vectors, added {added_num} vectors." diff --git a/hugegraph-llm/src/hugegraph_llm/flows/utils.py b/hugegraph-llm/src/hugegraph_llm/flows/utils.py new file mode 100644 index 000000000..b4ba05c84 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/utils.py @@ -0,0 +1,34 @@ +# 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 hugegraph_llm.state.ai_state import WkFlowInput +from hugegraph_llm.utils.log import log + + +def prepare_schema(prepared_input: WkFlowInput, schema): + schema = schema.strip() + if schema.startswith("{"): + try: + schema = json.loads(schema) + prepared_input.schema = schema + except json.JSONDecodeError as exc: + log.error("Invalid JSON format in schema. Please check it again.") + raise ValueError("Invalid JSON format in schema.") from exc + else: + log.info("Get schema '%s' from graphdb.", schema) + prepared_input.graph_name = schema + return diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py new file mode 100644 index 000000000..0ea0675c0 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py @@ -0,0 +1,71 @@ +# 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 PyCGraph import GNode, CStatus +from hugegraph_llm.nodes.util import init_context +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class BaseNode(GNode): + context: WkFlowState = None + wk_input: WkFlowInput = None + + def init(self): + return init_context(self) + + def node_init(self): + """ + Node initialization method, can be overridden by subclasses. + Returns a CStatus object indicating whether initialization succeeded. + """ + return CStatus() + + def run(self): + """ + Main logic for node execution, can be overridden by subclasses. + Returns a CStatus object indicating whether execution succeeded. + """ + sts = self.node_init() + if sts.isErr(): + return sts + self.context.lock() + try: + data_json = self.context.to_json() + finally: + self.context.unlock() + + try: + res = self.operator_schedule(data_json) + except Exception as exc: + import traceback + + node_info = f"Node type: {type(self).__name__}, Node object: {self}" + err_msg = f"Node failed: {exc}\n{node_info}\n{traceback.format_exc()}" + return CStatus(-1, err_msg) + + self.context.lock() + try: + if isinstance(res, dict): + self.context.assign_from_json(res) + finally: + self.context.unlock() + return CStatus() + + def operator_schedule(self, data_json): + """ + Interface for scheduling the operator, can be overridden by subclasses. + Returns a CStatus object indicating whether scheduling succeeded. + """ + pass 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 new file mode 100644 index 000000000..4c5acbe97 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/document_node/chunk_split.py @@ -0,0 +1,43 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.operators.document_op.chunk_split import ChunkSplit +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class ChunkSplitNode(BaseNode): + chunk_split_op: ChunkSplit + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + if ( + self.wk_input.texts is None + or self.wk_input.language is None + or self.wk_input.split_type is None + ): + return CStatus(-1, "Error occurs when prepare for workflow input") + texts = self.wk_input.texts + language = self.wk_input.language + split_type = self.wk_input.split_type + if isinstance(texts, str): + texts = [texts] + self.chunk_split_op = ChunkSplit(texts, split_type, language) + return CStatus() + + def operator_schedule(self, data_json): + 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 new file mode 100644 index 000000000..b576e8170 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/commit_to_hugegraph.py @@ -0,0 +1,35 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph import Commit2Graph +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class Commit2GraphNode(BaseNode): + commit_to_graph_op: Commit2Graph + context: WkFlowState = None + wk_input: WkFlowInput = None + + 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() + return CStatus() + + def operator_schedule(self, data_json): + return self.commit_to_graph_op.run(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 new file mode 100644 index 000000000..b2434e524 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/fetch_graph_data.py @@ -0,0 +1,33 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.hugegraph_op.fetch_graph_data import FetchGraphData +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.utils.hugegraph_utils import get_hg_client + + +class FetchGraphDataNode(BaseNode): + fetch_graph_data_op: FetchGraphData + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + self.fetch_graph_data_op = FetchGraphData(get_hg_client()) + return CStatus() + + def operator_schedule(self, data_json): + return self.fetch_graph_data_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py new file mode 100644 index 000000000..71c490b20 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.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 json + +from PyCGraph import CStatus +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.common_op.check_schema import CheckSchema +from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.utils.log import log + + +class SchemaNode(BaseNode): + schema_manager: SchemaManager + check_schema: CheckSchema + context: WkFlowState = None + wk_input: WkFlowInput = None + + schema = None + + def _import_schema( + self, + from_hugegraph=None, + from_extraction=None, + from_user_defined=None, + ): + if from_hugegraph: + return SchemaManager(from_hugegraph) + elif from_user_defined: + return CheckSchema(from_user_defined) + elif from_extraction: + raise NotImplementedError("Not implemented yet") + else: + raise ValueError("No input data / invalid schema type") + + def node_init(self): + self.schema = self.wk_input.schema + self.schema = self.schema.strip() + if self.schema.startswith("{"): + try: + schema = json.loads(self.schema) + self.check_schema = self._import_schema(from_user_defined=schema) + except json.JSONDecodeError as exc: + log.error("Invalid JSON format in schema. Please check it again.") + raise ValueError("Invalid JSON format in schema.") from exc + else: + log.info("Get schema '%s' from graphdb.", self.schema) + self.schema_manager = self._import_schema(from_hugegraph=self.schema) + return CStatus() + + def operator_schedule(self, data_json): + print(f"check data json {data_json}") + if self.schema.startswith("{"): + try: + return self.check_schema.run(data_json) + except json.JSONDecodeError as exc: + log.error("Invalid JSON format in schema. Please check it again.") + raise ValueError("Invalid JSON format in schema.") from exc + else: + log.info("Get schema '%s' from graphdb.", self.schema) + return self.schema_manager.run(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 new file mode 100644 index 000000000..ab31fa394 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_semantic_index.py @@ -0,0 +1,34 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.config import llm_settings +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class BuildSemanticIndexNode(BaseNode): + build_semantic_index_op: BuildSemanticIndex + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + self.build_semantic_index_op = BuildSemanticIndex(get_embedding(llm_settings)) + return CStatus() + + def operator_schedule(self, data_json): + return self.build_semantic_index_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py new file mode 100644 index 000000000..cf2f9b677 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py @@ -0,0 +1,34 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.config import llm_settings +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class BuildVectorIndexNode(BaseNode): + build_vector_index_op: BuildVectorIndex + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + self.build_vector_index_op = BuildVectorIndex(get_embedding(llm_settings)) + return CStatus() + + def operator_schedule(self, data_json): + return self.build_vector_index_op.run(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 new file mode 100644 index 000000000..8bceed804 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/extract_info.py @@ -0,0 +1,52 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.config import llm_settings +from hugegraph_llm.models.llms.init_llm import get_chat_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 +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class ExtractNode(BaseNode): + property_graph_extract: PropertyGraphExtract + info_extract: InfoExtract + context: WkFlowState = None + wk_input: WkFlowInput = None + + extract_type: str = None + + def node_init(self): + llm = get_chat_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 + extract_type = self.wk_input.extract_type + self.extract_type = extract_type + if extract_type == "triples": + self.info_extract = InfoExtract(llm, example_prompt) + elif extract_type == "property_graph": + self.property_graph_extract = PropertyGraphExtract(llm, example_prompt) + else: + return CStatus(-1, f"Unsupported extract_type: {extract_type}") + return CStatus() + + def operator_schedule(self, data_json): + if self.extract_type == "triples": + return self.info_extract.run(data_json) + elif self.extract_type == "property_graph": + return self.property_graph_extract.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py new file mode 100644 index 000000000..317f9e6ac --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py @@ -0,0 +1,59 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.config import llm_settings +from hugegraph_llm.models.llms.init_llm import get_chat_llm +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.llm_op.prompt_generate import PromptGenerate +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class PromptGenerateNode(BaseNode): + prompt_generate: PromptGenerate + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + """ + Node initialization method, initialize PromptGenerate operator + """ + llm = get_chat_llm(llm_settings) + if not all( + [ + self.wk_input.source_text, + self.wk_input.scenario, + self.wk_input.example_name, + ] + ): + return CStatus( + -1, + "Missing required parameters: source_text, scenario, or example_name", + ) + + self.prompt_generate = PromptGenerate(llm) + context = { + "source_text": self.wk_input.source_text, + "scenario": self.wk_input.scenario, + "example_name": self.wk_input.example_name, + } + self.context.assign_from_json(context) + return CStatus() + + def operator_schedule(self, data_json): + """ + Schedule the execution of PromptGenerate operator + """ + return self.prompt_generate.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py new file mode 100644 index 000000000..a28b41346 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py @@ -0,0 +1,91 @@ +# 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 PyCGraph import CStatus +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.models.llms.init_llm import get_chat_llm +from hugegraph_llm.config import llm_settings +from hugegraph_llm.operators.llm_op.schema_build import SchemaBuilder +from hugegraph_llm.utils.log import log + + +class SchemaBuildNode(BaseNode): + schema_builder: SchemaBuilder + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + llm = get_chat_llm(llm_settings) + self.schema_builder = SchemaBuilder(llm) + + # texts -> raw_texts + raw_texts = [] + if self.wk_input.texts: + if isinstance(self.wk_input.texts, list): + raw_texts = [t for t in self.wk_input.texts if isinstance(t, str)] + elif isinstance(self.wk_input.texts, str): + raw_texts = [self.wk_input.texts] + + # query_examples: already parsed list[dict] or raw JSON string + query_examples = [] + qe_src = self.wk_input.query_examples if self.wk_input.query_examples else None + if qe_src: + try: + parsed_examples = json.loads(qe_src) + # Validate and retain the description and gremlin fields + query_examples = [ + { + "description": ex.get("description", ""), + "gremlin": ex.get("gremlin", ""), + } + for ex in parsed_examples + if isinstance(ex, dict) and "description" in ex and "gremlin" in ex + ] + except json.JSONDecodeError as e: + return CStatus(-1, f"Query Examples is not in a valid JSON format: {e}") + + # few_shot_schema: already parsed dict or raw JSON string + few_shot_schema = {} + fss_src = ( + self.wk_input.few_shot_schema if self.wk_input.few_shot_schema else None + ) + if fss_src: + try: + few_shot_schema = json.loads(fss_src) + except json.JSONDecodeError as e: + return CStatus( + -1, f"Few Shot Schema is not in a valid JSON format: {e}" + ) + + _context_payload = { + "raw_texts": raw_texts, + "query_examples": query_examples, + "few_shot_schema": few_shot_schema, + } + self.context.assign_from_json(_context_payload) + + return CStatus() + + def operator_schedule(self, data_json): + try: + schema_result = self.schema_builder.run(data_json) + + return {"schema": schema_result} + except Exception as e: + log.error("Failed to generate schema: %s", e) + return {"schema": f"Schema generation failed: {e}"} diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/util.py b/hugegraph-llm/src/hugegraph_llm/nodes/util.py new file mode 100644 index 000000000..60bdc2e86 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/util.py @@ -0,0 +1,27 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from PyCGraph import CStatus + + +def init_context(obj) -> CStatus: + try: + obj.context = obj.getGParamWithNoEmpty("wkflow_state") + obj.wk_input = obj.getGParamWithNoEmpty("wkflow_input") + if obj.context is None or obj.wk_input is None: + return CStatus(-1, "Required workflow parameters not found") + return CStatus() + except Exception as e: + return CStatus(-1, f"Failed to initialize context: {str(e)}") diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index 7a533517a..c1c742032 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -20,12 +20,8 @@ from hugegraph_llm.enums.property_cardinality import PropertyCardinality from hugegraph_llm.enums.property_data_type import PropertyDataType -from hugegraph_llm.operators.util import init_context from hugegraph_llm.utils.log import log -from PyCGraph import GNode, CStatus -from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState - def log_and_raise(message: str) -> None: log.warning(message) @@ -174,159 +170,3 @@ def _add_missing_properties( } ) property_label_set.add(prop) - - -class CheckSchemaNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None - - def init(self): - return init_context(self) - - def node_init(self): - if self.wk_input.schema is None: - return CStatus(-1, "Error occurs when prepare for workflow input") - self.data = self.wk_input.schema - return CStatus() - - def run(self) -> CStatus: - # init workflow input - sts = self.node_init() - if sts.isErr(): - return sts - # 1. Validate the schema structure - self.context.lock() - schema = self.data or self.context.schema - self._validate_schema(schema) - # 2. Process property labels and also create a set for it - property_labels, property_label_set = self._process_property_labels(schema) - # 3. Process properties in given vertex/edge labels - self._process_vertex_labels(schema, property_labels, property_label_set) - self._process_edge_labels(schema, property_labels, property_label_set) - # 4. Update schema with processed pks - schema["propertykeys"] = property_labels - self.context.schema = schema - self.context.unlock() - return CStatus() - - def _validate_schema(self, schema: Dict[str, Any]) -> None: - check_type(schema, dict, "Input data is not a dictionary.") - if "vertexlabels" not in schema or "edgelabels" not in schema: - log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") - check_type( - schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." - ) - check_type( - schema["edgelabels"], list, "'edgelabels' in input data is not a list." - ) - - def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): - property_labels = schema.get("propertykeys", []) - check_type( - property_labels, - list, - "'propertykeys' in input data is not of correct type.", - ) - property_label_set = {label["name"] for label in property_labels} - return property_labels, property_label_set - - def _process_vertex_labels( - self, schema: Dict[str, Any], property_labels: list, property_label_set: set - ) -> None: - for vertex_label in schema["vertexlabels"]: - self._validate_vertex_label(vertex_label) - properties = vertex_label["properties"] - primary_keys = self._process_keys( - vertex_label, "primary_keys", properties[:1] - ) - if len(primary_keys) == 0: - log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") - vertex_label["primary_keys"] = primary_keys - nullable_keys = self._process_keys( - vertex_label, "nullable_keys", properties[1:] - ) - vertex_label["nullable_keys"] = nullable_keys - self._add_missing_properties( - properties, property_labels, property_label_set - ) - - def _process_edge_labels( - self, schema: Dict[str, Any], property_labels: list, property_label_set: set - ) -> None: - for edge_label in schema["edgelabels"]: - self._validate_edge_label(edge_label) - properties = edge_label.get("properties", []) - self._add_missing_properties( - properties, property_labels, property_label_set - ) - - def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: - check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") - if "name" not in vertex_label: - log_and_raise("VertexLabel in input data does not contain 'name'.") - check_type( - vertex_label["name"], str, "'name' in vertex_label is not of correct type." - ) - if "properties" not in vertex_label: - log_and_raise("VertexLabel in input data does not contain 'properties'.") - check_type( - vertex_label["properties"], - list, - "'properties' in vertex_label is not of correct type.", - ) - if len(vertex_label["properties"]) == 0: - log_and_raise("'properties' in vertex_label is empty.") - - def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: - check_type(edge_label, dict, "EdgeLabel in input data is not a dictionary.") - if ( - "name" not in edge_label - or "source_label" not in edge_label - or "target_label" not in edge_label - ): - log_and_raise( - "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." - ) - check_type( - edge_label["name"], str, "'name' in edge_label is not of correct type." - ) - check_type( - edge_label["source_label"], - str, - "'source_label' in edge_label is not of correct type.", - ) - check_type( - edge_label["target_label"], - str, - "'target_label' in edge_label is not of correct type.", - ) - - def _process_keys( - self, label: Dict[str, Any], key_type: str, default_keys: list - ) -> list: - keys = label.get(key_type, default_keys) - check_type( - keys, list, f"'{key_type}' in {label['name']} is not of correct type." - ) - new_keys = [key for key in keys if key in label["properties"]] - return new_keys - - def _add_missing_properties( - self, properties: list, property_labels: list, property_label_set: set - ) -> None: - for prop in properties: - if prop not in property_label_set: - property_labels.append( - { - "name": prop, - "data_type": PropertyDataType.DEFAULT.value, - "cardinality": PropertyCardinality.DEFAULT.value, - } - ) - property_label_set.add(prop) - - def get_result(self): - self.context.lock() - res = self.context.to_json() - self.context.unlock() - return res 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 d779a40ab..c31e77af7 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 @@ -19,8 +19,6 @@ from typing import Literal, Dict, Any, Optional, Union, List from langchain_text_splitters import RecursiveCharacterTextSplitter -from hugegraph_llm.operators.util import init_context -from PyCGraph import GNode, CStatus # Constants LANGUAGE_ZH = "zh" @@ -30,62 +28,6 @@ SPLIT_TYPE_SENTENCE = "sentence" -class ChunkSplitNode(GNode): - def init(self): - return init_context(self) - - def node_init(self): - if ( - self.wk_input.texts is None - or self.wk_input.language is None - or self.wk_input.split_type is None - ): - return CStatus(-1, "Error occurs when prepare for workflow input") - texts = self.wk_input.texts - language = self.wk_input.language - split_type = self.wk_input.split_type - if isinstance(texts, str): - texts = [texts] - self.texts = texts - self.separators = self._get_separators(language) - self.text_splitter = self._get_text_splitter(split_type) - return CStatus() - - def _get_separators(self, language: str) -> List[str]: - if language == LANGUAGE_ZH: - return ["\n\n", "\n", "。", "，", ""] - if language == LANGUAGE_EN: - return ["\n\n", "\n", ".", ",", " ", ""] - raise ValueError("language must be zh or en") - - 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 - if split_type == SPLIT_TYPE_SENTENCE: - return RecursiveCharacterTextSplitter( - chunk_size=50, chunk_overlap=0, separators=self.separators - ).split_text - raise ValueError("Type must be document, paragraph or sentence") - - def run(self): - sts = self.node_init() - if sts.isErr(): - return sts - all_chunks = [] - for text in self.texts: - chunks = self.text_splitter(text) - all_chunks.extend(chunks) - - self.context.lock() - self.context.chunks = all_chunks - self.context.unlock() - return CStatus() - - class ChunkSplit: def __init__( self, 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 5cc846d21..9eec04f7f 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 @@ -40,15 +40,19 @@ def run(self, data: dict) -> Dict[str, Any]: schema = data.get("schema") vertices = data.get("vertices", []) edges = data.get("edges", []) - + print(f"get schema {schema}") if not vertices and not edges: - log.critical("(Loading) Both vertices and edges are empty. Please check the input data again.") + log.critical( + "(Loading) Both vertices and edges are empty. Please check the input data again." + ) raise ValueError("Both vertices and edges input are empty.") if not schema: # TODO: ensure the function works correctly (update the logic later) self.schema_free_mode(data.get("triples", [])) - log.warning("Using schema_free mode, could try schema_define mode for better effect!") + log.warning( + "Using schema_free mode, could try schema_define mode for better effect!" + ) else: self.init_schema_if_need(schema) self.load_into_graph(vertices, edges, schema) @@ -64,7 +68,9 @@ def _set_default_property(self, key, input_properties, property_label_map): # list or set default_value = [] input_properties[key] = default_value - log.warning("Property '%s' missing in vertex, set to '%s' for now", key, default_value) + log.warning( + "Property '%s' missing in vertex, set to '%s' for now", key, default_value + ) def _handle_graph_creation(self, func, *args, **kwargs): try: @@ -78,29 +84,42 @@ 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) - vertex_label_map = {v_label["name"]: v_label for v_label in schema["vertexlabels"]} + 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"]} + property_label_map = { + p_label["name"]: p_label for p_label in schema["propertykeys"] + } for vertex in vertices: input_label = vertex["label"] # 1. ensure the input_label in the graph schema if input_label not in vertex_label_map: - log.critical("(Input) VertexLabel %s not found in schema, skip & need check it!", input_label) + log.critical( + "(Input) VertexLabel %s not found in schema, skip & need check it!", + input_label, + ) continue input_properties = vertex["properties"] vertex_label = vertex_label_map[input_label] primary_keys = vertex_label["primary_keys"] nullable_keys = vertex_label.get("nullable_keys", []) - non_null_keys = [key for key in vertex_label["properties"] if key not in nullable_keys] + non_null_keys = [ + key for key in vertex_label["properties"] if key not in nullable_keys + ] has_problem = False # 2. Handle primary-keys mode vertex for pk in primary_keys: if not input_properties.get(pk): if len(primary_keys) == 1: - log.error("Primary-key '%s' missing in vertex %s, skip it & need check it again", pk, vertex) + log.error( + "Primary-key '%s' missing in vertex %s, skip it & need check it again", + pk, + vertex, + ) has_problem = True break # TODO: transform to Enum first (better in earlier step) @@ -110,14 +129,20 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- input_properties[pk] = default_value_map(data_type) else: input_properties[pk] = [] - log.warning("Primary-key '%s' missing in vertex %s, mark empty & need check it again!", pk, vertex) + log.warning( + "Primary-key '%s' missing in vertex %s, mark empty & need check it again!", + pk, + vertex, + ) if has_problem: continue # 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) + 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(): @@ -125,14 +150,19 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- 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) + log.error( + "Property type/format '%s' is not correct, skip it & need check it again", + key, + ) has_problem = True break if has_problem: continue # TODO: we could try batch add vertices first, setback to single-mode if failed - vid = self._handle_graph_creation(self.client.graph().addVertex, input_label, input_properties).id + vid = self._handle_graph_creation( + self.client.graph().addVertex, input_label, input_properties + ).id vertex["id"] = vid for edge in edges: @@ -142,11 +172,16 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- properties = edge["properties"] if label not in edge_label_map: - log.critical("(Input) EdgeLabel %s not found in schema, skip & need check it!", label) + log.critical( + "(Input) EdgeLabel %s not found in schema, skip & need check it!", + label, + ) 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) + self._handle_graph_creation( + self.client.graph().addEdge, label, start, end, properties + ) def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -170,19 +205,27 @@ def init_schema_if_need(self, schema: dict): source_vertex_label = edge["source_label"] target_vertex_label = edge["target_label"] properties = edge["properties"] - self.schema.edgeLabel(edge_label).sourceLabel(source_vertex_label).targetLabel( - target_vertex_label - ).properties(*properties).nullableKeys(*properties).ifNotExist().create() + self.schema.edgeLabel(edge_label).sourceLabel( + source_vertex_label + ).targetLabel(target_vertex_label).properties(*properties).nullableKeys( + *properties + ).ifNotExist().create() def schema_free_mode(self, 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( + self.schema.vertexLabel("vertex").useCustomizeStringId().properties( "name" ).ifNotExist().create() + self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel( + "vertex" + ).properties("name").ifNotExist().create() - self.schema.indexLabel("vertexByName").onV("vertex").by("name").secondary().ifNotExist().create() - self.schema.indexLabel("edgeByName").onE("edge").by("name").secondary().ifNotExist().create() + 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: s, p, o = (element.strip() for element in item) @@ -196,8 +239,12 @@ def _create_property(self, prop: dict): data_type = PropertyDataType(prop["data_type"]) cardinality = PropertyCardinality(prop["cardinality"]) except ValueError: - log.critical("Invalid data type %s / cardinality %s for property %s, skip & should check it again", - prop["data_type"], prop["cardinality"], name) + log.critical( + "Invalid data type %s / cardinality %s for property %s, skip & should check it again", + prop["data_type"], + prop["cardinality"], + name, + ) return property_key = self.schema.propertyKey(name) @@ -231,7 +278,9 @@ def _set_property_data_type(self, property_key, data_type): log.warning("UUID type is not supported, use text instead") property_key.asText() else: - log.error("Unknown data type %s for property_key %s", data_type, property_key) + log.error( + "Unknown data type %s for property_key %s", data_type, property_key + ) def _set_property_cardinality(self, property_key, cardinality): if cardinality == PropertyCardinality.SINGLE: @@ -241,10 +290,17 @@ def _set_property_cardinality(self, property_key, cardinality): elif cardinality == PropertyCardinality.SET: property_key.valueSet() else: - 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): + 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) @@ -259,14 +315,21 @@ def _check_collection_data_type(self, data_type: str, value) -> bool: def _check_single_data_type(self, data_type: str, value) -> bool: if data_type == PropertyDataType.BOOLEAN.value: return isinstance(value, bool) - if data_type in (PropertyDataType.BYTE.value, PropertyDataType.INT.value, PropertyDataType.LONG.value): + if data_type in ( + PropertyDataType.BYTE.value, + PropertyDataType.INT.value, + PropertyDataType.LONG.value, + ): return isinstance(value, int) if data_type in (PropertyDataType.FLOAT.value, PropertyDataType.DOUBLE.value): return isinstance(value, float) if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below - if data_type == PropertyDataType.DATE.value: # the format should be "yyyy-MM-dd" + if ( + data_type == PropertyDataType.DATE.value + ): # the format should be "yyyy-MM-dd" import re - return isinstance(value, str) and re.match(r'^\d{4}-\d{2}-\d{2}$', value) + + return isinstance(value, str) and re.match(r"^\d{4}-\d{2}-\d{2}$", value) raise ValueError(f"Unknown/Unsupported data type: {data_type}") 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 670c18b4a..c4e2124c3 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 @@ -17,12 +17,8 @@ from typing import Dict, Any, Optional from hugegraph_llm.config import huge_settings -from hugegraph_llm.operators.util import init_context -from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState from pyhugegraph.client import PyHugeClient -from PyCGraph import GNode, CStatus - class SchemaManager: def __init__(self, graph_name: str): @@ -74,74 +70,3 @@ def run(self, context: Optional[Dict[str, Any]]) -> Dict[str, Any]: # TODO: enhance the logic here context["simple_schema"] = self.simple_schema(schema) return context - - -class SchemaManagerNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None - - def init(self): - return init_context(self) - - def node_init(self): - if self.wk_input.graph_name is None: - return CStatus(-1, "Error occurs when prepare for workflow input") - graph_name = self.wk_input.graph_name - self.graph_name = graph_name - self.client = PyHugeClient( - url=huge_settings.graph_url, - graph=self.graph_name, - user=huge_settings.graph_user, - pwd=huge_settings.graph_pwd, - graphspace=huge_settings.graph_space, - ) - self.schema = self.client.schema() - return CStatus() - - def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: - mini_schema = {} - - # Add necessary vertexlabels items (3) - if "vertexlabels" in schema: - mini_schema["vertexlabels"] = [] - for vertex in schema["vertexlabels"]: - new_vertex = { - key: vertex[key] - for key in ["id", "name", "properties"] - if key in vertex - } - mini_schema["vertexlabels"].append(new_vertex) - - # Add necessary edgelabels items (4) - if "edgelabels" in schema: - mini_schema["edgelabels"] = [] - for edge in schema["edgelabels"]: - new_edge = { - key: edge[key] - for key in ["name", "source_label", "target_label", "properties"] - if key in edge - } - mini_schema["edgelabels"].append(new_edge) - - return mini_schema - - def run(self) -> CStatus: - sts = self.node_init() - if sts.isErr(): - return sts - schema = self.schema.getSchema() - if not schema["vertexlabels"] and not schema["edgelabels"]: - raise Exception(f"Can not get {self.graph_name}'s schema from HugeGraph!") - - self.context.lock() - self.context.schema = schema - # TODO: enhance the logic here - self.context.simple_schema = self.simple_schema(schema) - self.context.unlock() - return CStatus() - - def get_result(self): - self.context.lock() - res = self.context.to_json() - self.context.unlock() - return res diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py index ee89d330f..5cdad0316 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py @@ -30,54 +30,6 @@ ) from hugegraph_llm.utils.log import log -from hugegraph_llm.operators.util import init_context -from hugegraph_llm.models.embeddings.init_embedding import get_embedding -from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState -from PyCGraph import GNode, CStatus - - -class BuildVectorIndexNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None - - def init(self): - return init_context(self) - - def node_init(self): - self.embedding = get_embedding(llm_settings) - self.folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) - self.filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(self.embedding, "model_name", None) - ) - self.vector_index = VectorIndex.from_index_file( - self.index_dir, self.filename_prefix - ) - return CStatus() - - def run(self): - # init workflow input - sts = self.node_init() - if sts.isErr(): - return sts - self.context.lock() - try: - if self.context.chunks is None: - raise ValueError("chunks not found in context.") - chunks = self.context.chunks - finally: - self.context.unlock() - chunks_embedding = [] - log.debug("Building vector index for %s chunks...", len(chunks)) - # TODO: use async_get_texts_embedding instead of single sync method - chunks_embedding = asyncio.run(get_embeddings_parallel(self.embedding, chunks)) - if len(chunks_embedding) > 0: - self.vector_index.add(chunks_embedding, chunks) - self.vector_index.to_index_file(self.index_dir, self.filename_prefix) - return CStatus() - class BuildVectorIndex: def __init__(self, embedding: BaseEmbedding): 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 15a8fdda7..571ffde51 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 @@ -18,16 +18,10 @@ import re from typing import List, Any, Dict, Optional -from hugegraph_llm.config import llm_settings from hugegraph_llm.document.chunk_split import ChunkSplitter from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.utils.log import log -from hugegraph_llm.operators.util import init_context -from hugegraph_llm.models.llms.init_llm import get_chat_llm -from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState -from PyCGraph import GNode, CStatus - SCHEMA_EXAMPLE_PROMPT = """## Main Task Extract Triples from the given text and graph schema @@ -213,143 +207,3 @@ def _filter_long_id(self, graph) -> Dict[str, List[Any]]: if self.valid(edge["start"]) and self.valid(edge["end"]) ] return graph - - -class InfoExtractNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None - - def init(self): - return init_context(self) - - def node_init(self): - self.llm = get_chat_llm(llm_settings) - if self.wk_input.example_prompt is None: - return CStatus(-1, "Error occurs when prepare for workflow input") - self.example_prompt = self.wk_input.example_prompt - return CStatus() - - def extract_triples_by_regex_with_schema(self, schema, text): - text = text.replace("\\n", " ").replace("\\", " ").replace("\n", " ") - pattern = r"\((.*?), (.*?), (.*?)\) - ([^ ]*)" - matches = re.findall(pattern, text) - - vertices_dict = {v["id"]: v for v in self.context.vertices} - for match in matches: - s, p, o, label = [item.strip() for item in match] - if None in [label, s, p, o]: - continue - # TODO: use a more efficient way to compare the extract & input property - p_lower = p.lower() - for vertex in schema["vertices"]: - if vertex["vertex_label"] == label and any( - pp.lower() == p_lower for pp in vertex["properties"] - ): - id = f"{label}-{s}" - if id not in vertices_dict: - vertices_dict[id] = { - "id": id, - "name": s, - "label": label, - "properties": {p: o}, - } - else: - vertices_dict[id]["properties"].update({p: o}) - break - for edge in schema["edges"]: - if edge["edge_label"] == label: - source_label = edge["source_vertex_label"] - source_id = f"{source_label}-{s}" - if source_id not in vertices_dict: - vertices_dict[source_id] = { - "id": source_id, - "name": s, - "label": source_label, - "properties": {}, - } - target_label = edge["target_vertex_label"] - target_id = f"{target_label}-{o}" - if target_id not in vertices_dict: - vertices_dict[target_id] = { - "id": target_id, - "name": o, - "label": target_label, - "properties": {}, - } - self.context.edges.append( - { - "start": source_id, - "end": target_id, - "type": label, - "properties": {}, - } - ) - break - self.context.vertices = list(vertices_dict.values()) - - def extract_triples_by_regex(self, text): - text = text.replace("\\n", " ").replace("\\", " ").replace("\n", " ") - pattern = r"\((.*?), (.*?), (.*?)\)" - self.context.triples += re.findall(pattern, text) - - def run(self) -> CStatus: - sts = self.node_init() - if sts.isErr(): - return sts - self.context.lock() - if self.context.chunks is None: - self.context.unlock() - raise ValueError("parameter required by extract node not found in context.") - schema = self.context.schema - chunks = self.context.chunks - - if schema: - self.context.vertices = [] - self.context.edges = [] - else: - self.context.triples = [] - - self.context.unlock() - - for sentence in chunks: - proceeded_chunk = self.extract_triples_by_llm(schema, sentence) - log.debug( - "[Legacy] %s input: %s \n output:%s", - self.__class__.__name__, - sentence, - proceeded_chunk, - ) - if schema: - self.extract_triples_by_regex_with_schema(schema, proceeded_chunk) - else: - self.extract_triples_by_regex(proceeded_chunk) - - if self.context.call_count: - self.context.call_count += len(chunks) - else: - self.context.call_count = len(chunks) - self._filter_long_id() - return CStatus() - - def extract_triples_by_llm(self, schema, chunk) -> str: - prompt = generate_extract_triple_prompt(chunk, schema) - if self.example_prompt is not None: - prompt = self.example_prompt + prompt - return self.llm.generate(prompt=prompt) - - # TODO: make 'max_length' be a configurable param in settings.py/settings.cfg - def valid(self, element_id: str, max_length: int = 256) -> bool: - if len(element_id.encode("utf-8")) >= max_length: - log.warning("Filter out GraphElementID too long: %s", element_id) - return False - return True - - def _filter_long_id(self): - self.context.vertices = [ - vertex for vertex in self.context.vertices if self.valid(vertex["id"]) - ] - self.context.edges = [ - edge - for edge in self.context.edges - if self.valid(edge["start"]) and self.valid(edge["end"]) - ] 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 6e492b8f5..79fb33b4f 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 @@ -21,16 +21,11 @@ import re from typing import List, Any, Dict -from hugegraph_llm.config import llm_settings, prompt +from hugegraph_llm.config import prompt from hugegraph_llm.document.chunk_split import ChunkSplitter from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.utils.log import log -from hugegraph_llm.operators.util import init_context -from hugegraph_llm.models.llms.init_llm import get_chat_llm -from hugegraph_llm.state.ai_state import WkFlowState, WkFlowInput -from PyCGraph import GNode, CStatus - # TODO: It is not clear whether there is any other dependence on the SCHEMA_EXAMPLE_PROMPT variable. # Because the SCHEMA_EXAMPLE_PROMPT variable will no longer change based on # prompt.extract_graph_prompt changes after the system loads, this does not seem to meet expectations. @@ -182,123 +177,3 @@ def process_items(item_list, valid_labels, item_type): "Invalid property graph JSON! Please check the extracted JSON data carefully" ) return items - - -class PropertyGraphExtractNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None - - def init(self): - self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name - return init_context(self) - - def node_init(self): - self.llm = get_chat_llm(llm_settings) - if self.wk_input.example_prompt is None: - return CStatus(-1, "Error occurs when prepare for workflow input") - self.example_prompt = self.wk_input.example_prompt - return CStatus() - - def run(self) -> CStatus: - sts = self.node_init() - if sts.isErr(): - return sts - self.context.lock() - try: - if self.context.schema is None or self.context.chunks is None: - raise ValueError( - "parameter required by extract node not found in context." - ) - schema = self.context.schema - chunks = self.context.chunks - if self.context.vertices is None: - self.context.vertices = [] - if self.context.edges is None: - self.context.edges = [] - finally: - self.context.unlock() - - items = [] - for chunk in chunks: - proceeded_chunk = self.extract_property_graph_by_llm(schema, chunk) - log.debug( - "[LLM] %s input: %s \n output:%s", - self.__class__.__name__, - chunk, - proceeded_chunk, - ) - items.extend(self._extract_and_filter_label(schema, proceeded_chunk)) - items = filter_item(schema, items) - self.context.lock() - try: - for item in items: - if item["type"] == "vertex": - self.context.vertices.append(item) - elif item["type"] == "edge": - self.context.edges.append(item) - finally: - self.context.unlock() - self.context.call_count = (self.context.call_count or 0) + len(chunks) - return CStatus() - - def extract_property_graph_by_llm(self, schema, chunk): - prompt = generate_extract_property_graph_prompt(chunk, schema) - if self.example_prompt is not None: - prompt = self.example_prompt + prompt - return self.llm.generate(prompt=prompt) - - def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: - # Use regex to extract a JSON object with curly braces - json_match = re.search(r"({.*})", text, re.DOTALL) - if not json_match: - log.critical( - "Invalid property graph! No JSON object found, " - "please check the output format example in prompt." - ) - return [] - json_str = json_match.group(1).strip() - - items = [] - try: - property_graph = json.loads(json_str) - # 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'." - ) - return items - - # Create sets for valid vertex and edge labels based on the schema - vertex_label_set = {vertex["name"] for vertex in schema["vertexlabels"]} - edge_label_set = {edge["name"] for edge in schema["edgelabels"]} - - def process_items(item_list, valid_labels, item_type): - for item in item_list: - if not isinstance(item, dict): - log.warning( - "Invalid property graph item type '%s'.", type(item) - ) - continue - if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): - log.warning("Invalid item keys '%s'.", item.keys()) - continue - if item["label"] not in valid_labels: - log.warning( - "Invalid %s label '%s' has been ignored.", - item_type, - item["label"], - ) - continue - items.append(item) - - process_items(property_graph["vertices"], vertex_label_set, "vertex") - process_items(property_graph["edges"], edge_label_set, "edge") - except json.JSONDecodeError: - log.critical( - "Invalid property graph JSON! Please check the extracted JSON data carefully" - ) - return items diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 0543aa2b4..6d3418c00 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -25,6 +25,14 @@ class WkFlowInput(GParam): example_prompt: str = None # need by graph information extract schema: str = None # Schema information requeired by SchemaNode graph_name: str = None + data_json = None + extract_type = None + query_examples = None + few_shot_schema = None + # Fields related to PromptGenerate + source_text: str = None # Original text + scenario: str = None # Scenario description + example_name: str = None # Example name def reset(self, _: CStatus) -> None: self.texts = None @@ -33,6 +41,14 @@ def reset(self, _: CStatus) -> None: self.example_prompt = None self.schema = None self.graph_name = None + self.data_json = None + self.extract_type = None + self.query_examples = None + self.few_shot_schema = None + # PromptGenerate related configuration + self.source_text = None + self.scenario = None + self.example_name = None class WkFlowState(GParam): @@ -49,6 +65,8 @@ class WkFlowState(GParam): graph_result = None keywords_embeddings = None + generated_extract_prompt: Optional[str] = None + def setup(self): self.schema = None self.simple_schema = None @@ -63,6 +81,8 @@ def setup(self): self.graph_result = None self.keywords_embeddings = None + self.generated_extract_prompt = None + return CStatus() def to_json(self): @@ -79,3 +99,11 @@ def to_json(self): for k, v in self.__dict__.items() if not k.startswith("_") and v is not None } + + # Implement a method that assigns keys from data_json as WkFlowState member variables + def assign_from_json(self, data_json: dict): + """ + Assigns each key in the input json object as a member variable of WkFlowState. + """ + for k, v in data_json.items(): + setattr(self, k, v) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index f61b5f843..ccace69f2 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -36,6 +36,15 @@ def get_graph_index_info(): + try: + scheduler = SchedulerSingleton.get_instance() + return scheduler.schedule_flow("get_graph_index_info") + except Exception as e: # pylint: disable=broad-exception-caught + log.error(e) + raise gr.Error(str(e)) + + +def get_graph_index_info_old(): builder = KgBuilder( LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() ) @@ -150,6 +159,15 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: def update_vid_embedding(): + scheduler = SchedulerSingleton.get_instance() + try: + return scheduler.schedule_flow("update_vid_embeddings") + except Exception as e: # pylint: disable=broad-exception-caught + log.error(e) + raise gr.Error(str(e)) + + +def update_vid_embedding_old(): builder = KgBuilder( LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() ) @@ -166,6 +184,18 @@ def update_vid_embedding(): def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: + try: + scheduler = SchedulerSingleton.get_instance() + return scheduler.schedule_flow("import_graph_data", data, schema) + except Exception as e: # pylint: disable=W0718 + log.error(e) + traceback.print_exc() + # Note: can't use gr.Error here + gr.Warning(str(e) + " Please check the graph data format/type carefully.") + return data + + +def import_graph_data_old(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: data_json = json.loads(data.strip()) log.debug("Import graph data: %s", data) @@ -190,6 +220,16 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: def build_schema(input_text, query_example, few_shot): + scheduler = SchedulerSingleton.get_instance() + try: + return scheduler.schedule_flow( + "build_schema", input_text, query_example, few_shot + ) + except (TypeError, ValueError) as e: + raise gr.Error(f"Schema generation failed: {e}") + + +def build_schema_old(input_text, query_example, few_shot): context = { "raw_texts": [input_text] if input_text else [], "query_examples": [], From d36d41d164ed04b77bf0d105ced442638bb5a2be Mon Sep 17 00:00:00 2001 From: LingXiao Qi Date: Tue, 30 Sep 2025 00:24:18 +0800 Subject: [PATCH 58/71] refactor: text2germlin with PCGraph framework (#50) Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Linyu <94553312+weijinglin@users.noreply.github.com> --- .../src/hugegraph_llm/api/admin_api.py | 8 +- .../api/exceptions/rag_exceptions.py | 4 +- .../hugegraph_llm/api/models/rag_requests.py | 106 +++++++---- .../src/hugegraph_llm/api/rag_api.py | 43 +++-- .../src/hugegraph_llm/config/admin_config.py | 2 + .../src/hugegraph_llm/config/generate.py | 4 +- .../hugegraph_llm/config/hugegraph_config.py | 1 + .../src/hugegraph_llm/config/llm_config.py | 21 +- .../config/models/base_config.py | 22 ++- .../config/models/base_prompt_config.py | 24 ++- .../src/hugegraph_llm/config/prompt_config.py | 1 + .../demo/rag_demo/admin_block.py | 39 ++-- .../src/hugegraph_llm/demo/rag_demo/app.py | 4 +- .../demo/rag_demo/configs_block.py | 91 +++------ .../demo/rag_demo/other_block.py | 13 +- .../hugegraph_llm/demo/rag_demo/rag_block.py | 74 +++++--- .../demo/rag_demo/text2gremlin_block.py | 109 ++++++++--- .../demo/rag_demo/vector_graph_block.py | 74 +++----- .../src/hugegraph_llm/document/chunk_split.py | 14 +- .../flows/get_graph_index_info.py | 4 +- .../src/hugegraph_llm/flows/graph_extract.py | 4 +- .../hugegraph_llm/flows/import_graph_data.py | 8 +- .../hugegraph_llm/flows/prompt_generate.py | 4 +- .../src/hugegraph_llm/flows/scheduler.py | 9 +- .../src/hugegraph_llm/flows/text2gremlin.py | 112 +++++++++++ .../src/hugegraph_llm/indices/graph_index.py | 17 +- .../src/hugegraph_llm/indices/vector_index.py | 33 +++- .../hugegraph_llm/middleware/middleware.py | 5 +- .../hugegraph_llm/models/embeddings/base.py | 37 ++-- .../hugegraph_llm/models/embeddings/openai.py | 25 ++- .../src/hugegraph_llm/models/llms/base.py | 34 ++-- .../src/hugegraph_llm/models/llms/init_llm.py | 6 +- .../src/hugegraph_llm/models/llms/litellm.py | 10 +- .../src/hugegraph_llm/models/llms/ollama.py | 30 ++- .../src/hugegraph_llm/models/llms/openai.py | 4 +- .../hugegraph_llm/models/rerankers/cohere.py | 9 +- .../models/rerankers/init_reranker.py | 4 +- .../models/rerankers/siliconflow.py | 9 +- .../nodes/hugegraph_node/gremlin_execute.py | 68 +++++++ .../nodes/hugegraph_node/schema.py | 2 +- .../index_node/gremlin_example_index_query.py | 49 +++++ .../nodes/llm_node/schema_build.py | 8 +- .../nodes/llm_node/text2gremlin.py | 70 +++++++ .../operators/common_op/check_schema.py | 40 +--- .../operators/common_op/merge_dedup_rerank.py | 15 +- .../operators/document_op/word_extract.py | 3 +- .../hugegraph_llm/operators/graph_rag_task.py | 4 +- .../hugegraph_op/commit_to_hugegraph.py | 58 ++---- .../operators/hugegraph_op/graph_rag_query.py | 63 ++++-- .../operators/hugegraph_op/schema_manager.py | 4 +- .../index_op/build_gremlin_example_index.py | 14 +- .../index_op/build_semantic_index.py | 30 +-- .../operators/index_op/build_vector_index.py | 4 +- .../index_op/gremlin_example_index_query.py | 29 ++- .../operators/index_op/semantic_id_query.py | 33 ++-- .../operators/index_op/vector_index_query.py | 8 +- .../operators/kg_construction_task.py | 11 +- .../operators/llm_op/answer_synthesize.py | 179 ++++++++++++------ .../operators/llm_op/disambiguate_data.py | 3 +- .../operators/llm_op/gremlin_generate.py | 19 +- .../operators/llm_op/info_extract.py | 8 +- .../operators/llm_op/keyword_extract.py | 24 +++ .../operators/llm_op/prompt_generate.py | 6 +- .../llm_op/property_graph_extract.py | 18 +- .../operators/llm_op/schema_build.py | 14 +- .../src/hugegraph_llm/state/ai_state.py | 30 ++- .../src/hugegraph_llm/utils/anchor.py | 9 +- .../src/hugegraph_llm/utils/decorators.py | 1 + .../hugegraph_llm/utils/embedding_utils.py | 9 +- .../hugegraph_llm/utils/graph_index_utils.py | 40 +--- .../hugegraph_llm/utils/hugegraph_utils.py | 26 ++- hugegraph-llm/src/hugegraph_llm/utils/log.py | 2 +- .../hugegraph_llm/utils/vector_index_utils.py | 16 +- hugegraph-llm/src/tests/config/test_config.py | 1 + .../embeddings/test_openai_embedding.py | 1 + .../tests/models/llms/test_ollama_client.py | 7 +- .../operators/common_op/test_check_schema.py | 9 +- .../operators/common_op/test_nltk_helper.py | 1 + .../src/pyhugegraph/api/auth.py | 16 +- .../src/pyhugegraph/api/graph.py | 12 +- .../src/pyhugegraph/api/schema.py | 12 +- .../api/schema_manage/index_label.py | 12 +- .../src/pyhugegraph/api/services.py | 8 +- .../src/pyhugegraph/api/traverser.py | 45 ++--- .../src/pyhugegraph/client.py | 2 +- .../pyhugegraph/example/hugegraph_example.py | 14 +- .../structure/property_key_data.py | 4 +- .../src/pyhugegraph/utils/huge_config.py | 10 +- .../src/pyhugegraph/utils/huge_router.py | 8 +- .../src/pyhugegraph/utils/log.py | 4 +- .../src/pyhugegraph/utils/util.py | 23 ++- .../src/tests/api/test_auth.py | 8 +- .../src/tests/api/test_version.py | 8 +- .../src/tests/client_utils.py | 6 +- 94 files changed, 1330 insertions(+), 816 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/gremlin_execute.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py diff --git a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py index 05648d48e..4c192c29c 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py @@ -31,8 +31,12 @@ def admin_http_api(router: APIRouter, log_stream): @router.post("/logs", status_code=status.HTTP_200_OK) async def log_stream_api(req: LogStreamRequest): if admin_settings.admin_token != req.admin_token: - raise generate_response(RAGResponse(status_code=status.HTTP_403_FORBIDDEN, #pylint: disable=E0702 - message="Invalid admin_token")) + raise generate_response( + RAGResponse( + status_code=status.HTTP_403_FORBIDDEN, # pylint: disable=E0702 + message="Invalid admin_token", + ) + ) log_path = os.path.join("logs", req.log_file) # Create a StreamingResponse that reads from the log stream generator diff --git a/hugegraph-llm/src/hugegraph_llm/api/exceptions/rag_exceptions.py b/hugegraph-llm/src/hugegraph_llm/api/exceptions/rag_exceptions.py index 75eb14cf3..18723e30b 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/exceptions/rag_exceptions.py +++ b/hugegraph-llm/src/hugegraph_llm/api/exceptions/rag_exceptions.py @@ -21,7 +21,9 @@ class ExternalException(HTTPException): def __init__(self): - super().__init__(status_code=400, detail="Connect failed with error code -1, please check the input.") + super().__init__( + status_code=400, detail="Connect failed with error code -1, please check the input." + ) class ConnectionFailedException(HTTPException): diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py index cf227e8bd..f46aea02c 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py +++ b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py @@ -24,10 +24,10 @@ class GraphConfigRequest(BaseModel): - url: str = Query('127.0.0.1:8080', description="hugegraph client url.") - graph: str = Query('hugegraph', description="hugegraph client name.") - user: str = Query('', description="hugegraph client user.") - pwd: str = Query('', description="hugegraph client pwd.") + url: str = Query("127.0.0.1:8080", description="hugegraph client url.") + graph: str = Query("hugegraph", description="hugegraph client name.") + user: str = Query("", description="hugegraph client user.") + pwd: str = Query("", description="hugegraph client pwd.") gs: str = None @@ -36,22 +36,42 @@ class RAGRequest(BaseModel): raw_answer: bool = Query(False, description="Use LLM to generate answer directly") vector_only: bool = Query(False, description="Use LLM to generate answer with vector") graph_only: bool = Query(True, description="Use LLM to generate answer with graph RAG only") - graph_vector_answer: bool = Query(False, description="Use LLM to generate answer with vector & GraphRAG") + graph_vector_answer: bool = Query( + False, description="Use LLM to generate answer with vector & GraphRAG" + ) graph_ratio: float = Query(0.5, description="The ratio of GraphRAG ans & vector ans") - rerank_method: Literal["bleu", "reranker"] = Query("bleu", description="Method to rerank the results.") - near_neighbor_first: bool = Query(False, description="Prioritize near neighbors in the search results.") - custom_priority_info: str = Query("", description="Custom information to prioritize certain results.") + rerank_method: Literal["bleu", "reranker"] = Query( + "bleu", description="Method to rerank the results." + ) + near_neighbor_first: bool = Query( + False, description="Prioritize near neighbors in the search results." + ) + custom_priority_info: str = Query( + "", description="Custom information to prioritize certain results." + ) # Graph Configs - max_graph_items: int = Query(30, description="Maximum number of items for GQL queries in graph.") + max_graph_items: int = Query( + 30, description="Maximum number of items for GQL queries in graph." + ) topk_return_results: int = Query(20, description="Number of sorted results to return finally.") - vector_dis_threshold: float = Query(0.9, description="Threshold for vector similarity\ - (results greater than this will be ignored).") - topk_per_keyword: int = Query(1, description="TopK results returned for each keyword \ - extracted from the query, by default only the most similar one is returned.") - client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") + vector_dis_threshold: float = Query( + 0.9, + description="Threshold for vector similarity\ + (results greater than this will be ignored).", + ) + topk_per_keyword: int = Query( + 1, + description="TopK results returned for each keyword \ + extracted from the query, by default only the most similar one is returned.", + ) + client_config: Optional[GraphConfigRequest] = Query( + None, description="hugegraph server config." + ) # Keep prompt params in the end - answer_prompt: Optional[str] = Query(prompt.answer_prompt, description="Prompt to guide the answer generation.") + answer_prompt: Optional[str] = Query( + prompt.answer_prompt, description="Prompt to guide the answer generation." + ) keywords_extract_prompt: Optional[str] = Query( prompt.keywords_extract_prompt, description="Prompt for extracting keywords from query.", @@ -67,22 +87,39 @@ class RAGRequest(BaseModel): class GraphRAGRequest(BaseModel): query: str = Query(..., description="Query you want to ask") # Graph Configs - max_graph_items: int = Query(30, description="Maximum number of items for GQL queries in graph.") + max_graph_items: int = Query( + 30, description="Maximum number of items for GQL queries in graph." + ) topk_return_results: int = Query(20, description="Number of sorted results to return finally.") - vector_dis_threshold: float = Query(0.9, description="Threshold for vector similarity \ - (results greater than this will be ignored).") - topk_per_keyword: int = Query(1, description="TopK results returned for each keyword extracted\ - from the query, by default only the most similar one is returned.") + vector_dis_threshold: float = Query( + 0.9, + description="Threshold for vector similarity \ + (results greater than this will be ignored).", + ) + topk_per_keyword: int = Query( + 1, + description="TopK results returned for each keyword extracted\ + from the query, by default only the most similar one is returned.", + ) - client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") + client_config: Optional[GraphConfigRequest] = Query( + None, description="hugegraph server config." + ) get_vertex_only: bool = Query(False, description="return only keywords & vertex (early stop).") gremlin_tmpl_num: int = Query( - 1, description="Number of Gremlin templates to use. If num <=0 means template is not provided" + 1, + description="Number of Gremlin templates to use. If num <=0 means template is not provided", + ) + rerank_method: Literal["bleu", "reranker"] = Query( + "bleu", description="Method to rerank the results." + ) + near_neighbor_first: bool = Query( + False, description="Prioritize near neighbors in the search results." + ) + custom_priority_info: str = Query( + "", description="Custom information to prioritize certain results." ) - rerank_method: Literal["bleu", "reranker"] = Query("bleu", description="Method to rerank the results.") - near_neighbor_first: bool = Query(False, description="Prioritize near neighbors in the search results.") - custom_priority_info: str = Query("", description="Custom information to prioritize certain results.") gremlin_prompt: Optional[str] = Query( prompt.gremlin_generate_prompt, description="Prompt for the Text2Gremlin query.", @@ -115,6 +152,7 @@ class LogStreamRequest(BaseModel): admin_token: Optional[str] = None log_file: Optional[str] = "llm-server.log" + class GremlinOutputType(str, Enum): MATCH_RESULT = "match_result" TEMPLATE_GREMLIN = "template_gremlin" @@ -122,32 +160,36 @@ class GremlinOutputType(str, Enum): TEMPLATE_EXECUTION_RESULT = "template_execution_result" RAW_EXECUTION_RESULT = "raw_execution_result" + class GremlinGenerateRequest(BaseModel): query: str example_num: Optional[int] = Query( - 0, - description="Number of Gremlin templates to use.(0 means no templates)" + 0, description="Number of Gremlin templates to use.(0 means no templates)" ) gremlin_prompt: Optional[str] = Query( prompt.gremlin_generate_prompt, description="Prompt for the Text2Gremlin query.", ) - client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.") + client_config: Optional[GraphConfigRequest] = Query( + None, description="hugegraph server config." + ) output_types: Optional[List[GremlinOutputType]] = Query( default=[GremlinOutputType.TEMPLATE_GREMLIN], description=""" a list can contain "match_result","template_gremlin", "raw_gremlin","template_execution_result","raw_execution_result" You can specify which type of result do you need. Empty means all types. - """ + """, ) - @field_validator('gremlin_prompt') + @field_validator("gremlin_prompt") @classmethod def validate_prompt_placeholders(cls, v): if v is not None: - required_placeholders = ['{query}', '{schema}', '{example}', '{vertices}'] + required_placeholders = ["{query}", "{schema}", "{example}", "{vertices}"] missing = [p for p in required_placeholders if p not in v] if missing: - raise ValueError(f"Prompt template is missing required placeholders: {', '.join(missing)}") + raise ValueError( + f"Prompt template is missing required placeholders: {', '.join(missing)}" + ) return v diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py index bfa76e7ef..356176e4e 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py @@ -32,6 +32,8 @@ from hugegraph_llm.config import huge_settings from hugegraph_llm.config import llm_settings, prompt from hugegraph_llm.utils.log import log +from hugegraph_llm.flows.scheduler import SchedulerSingleton + # pylint: disable=too-many-statements @@ -74,7 +76,9 @@ def rag_answer_api(req: RAGRequest): "query": req.query, **{ key: value - for key, value in zip(["raw_answer", "vector_only", "graph_only", "graph_vector_answer"], result) + for key, value in zip( + ["raw_answer", "vector_only", "graph_only", "graph_vector_answer"], result + ) if getattr(req, key) }, } @@ -103,11 +107,12 @@ def graph_rag_recall_api(req: GraphRAGRequest): near_neighbor_first=req.near_neighbor_first, custom_related_information=req.custom_priority_info, gremlin_prompt=req.gremlin_prompt or prompt.gremlin_generate_prompt, - get_vertex_only=req.get_vertex_only + get_vertex_only=req.get_vertex_only, ) if req.get_vertex_only: from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery + graph_rag = GraphRAGQuery() graph_rag.init_client(result) vertex_details = graph_rag.get_vertex_details(result["match_vids"]) @@ -135,7 +140,8 @@ def graph_rag_recall_api(req: GraphRAGRequest): except Exception as e: log.error("Unexpected error occurred: %s", e) raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="An unexpected error occurred." + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="An unexpected error occurred.", ) from e @router.post("/config/graph", status_code=status.HTTP_201_CREATED) @@ -150,7 +156,9 @@ def llm_config_api(req: LLMConfigRequest): llm_settings.llm_type = req.llm_type if req.llm_type == "openai": - res = apply_llm_conf(req.api_key, req.api_base, req.language_model, req.max_tokens, origin_call="http") + res = apply_llm_conf( + req.api_key, req.api_base, req.language_model, req.max_tokens, origin_call="http" + ) else: res = apply_llm_conf(req.host, req.port, req.language_model, None, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) @@ -160,7 +168,9 @@ def embedding_config_api(req: LLMConfigRequest): llm_settings.embedding_type = req.llm_type if req.llm_type == "openai": - res = apply_embedding_conf(req.api_key, req.api_base, req.language_model, origin_call="http") + res = apply_embedding_conf( + req.api_key, req.api_base, req.language_model, origin_call="http" + ) else: res = apply_embedding_conf(req.host, req.port, req.language_model, origin_call="http") return generate_response(RAGResponse(status_code=res, message="Missing Value")) @@ -170,7 +180,9 @@ def rerank_config_api(req: RerankerConfigRequest): llm_settings.reranker_type = req.reranker_type if req.reranker_type == "cohere": - res = apply_reranker_conf(req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http") + res = apply_reranker_conf( + req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http" + ) elif req.reranker_type == "siliconflow": res = apply_reranker_conf(req.api_key, req.reranker_model, None, origin_call="http") else: @@ -182,16 +194,23 @@ def text2gremlin_api(req: GremlinGenerateRequest): try: set_graph_config(req) + # Basic parameter validation: empty query => 400 + if not req.query or not str(req.query).strip(): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail="Query must not be empty." + ) + output_types_str_list = None if req.output_types: output_types_str_list = [ot.value for ot in req.output_types] - response_dict = gremlin_generate_selective_func( - inp=req.query, - example_num=req.example_num, - schema_input=huge_settings.graph_name, - gremlin_prompt_input=req.gremlin_prompt, - requested_outputs=output_types_str_list, + response_dict = SchedulerSingleton.get_instance().schedule_flow( + "text2gremlin", + req.query, + req.example_num, + huge_settings.graph_name, + req.gremlin_prompt, + output_types_str_list, ) return response_dict except HTTPException as e: diff --git a/hugegraph-llm/src/hugegraph_llm/config/admin_config.py b/hugegraph-llm/src/hugegraph_llm/config/admin_config.py index b2814de41..fabc75de4 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/admin_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/admin_config.py @@ -18,8 +18,10 @@ from typing import Optional from .models import BaseConfig + class AdminConfig(BaseConfig): """Admin settings""" + enable_login: Optional[str] = "False" user_token: Optional[str] = "4321" admin_token: Optional[str] = "xxxx" diff --git a/hugegraph-llm/src/hugegraph_llm/config/generate.py b/hugegraph-llm/src/hugegraph_llm/config/generate.py index 36910e480..4b40e899f 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/generate.py +++ b/hugegraph-llm/src/hugegraph_llm/config/generate.py @@ -22,7 +22,9 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(description="Generate hugegraph-llm config file") - parser.add_argument("-U", "--update", default=True, action="store_true", help="Update the config file") + parser.add_argument( + "-U", "--update", default=True, action="store_true", help="Update the config file" + ) args = parser.parse_args() if args.update: huge_settings.generate_env() diff --git a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py b/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py index e51008d96..69abf0fbc 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py @@ -21,6 +21,7 @@ class HugeGraphConfig(BaseConfig): """HugeGraph settings""" + # graph server config graph_url: Optional[str] = "127.0.0.1:8080" graph_name: Optional[str] = "hugegraph" diff --git a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py index 64d851f5a..eb094ef88 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/llm_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/llm_config.py @@ -24,6 +24,7 @@ class LLMConfig(BaseConfig): """LLM settings""" + language: Literal["EN", "CN"] = "EN" chat_llm_type: Literal["openai", "litellm", "ollama/local"] = "openai" extract_llm_type: Literal["openai", "litellm", "ollama/local"] = "openai" @@ -35,23 +36,33 @@ class LLMConfig(BaseConfig): hybrid_llm_weights: Optional[float] = 0.5 # 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") + openai_chat_api_base: Optional[str] = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_chat_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") openai_chat_language_model: Optional[str] = "gpt-4.1-mini" - openai_extract_api_base: Optional[str] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_extract_api_base: Optional[str] = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_extract_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") openai_extract_language_model: Optional[str] = "gpt-4.1-mini" - openai_text2gql_api_base: Optional[str] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + openai_text2gql_api_base: Optional[str] = os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ) openai_text2gql_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") openai_text2gql_language_model: Optional[str] = "gpt-4.1-mini" - openai_embedding_api_base: Optional[str] = os.environ.get("OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1") + openai_embedding_api_base: Optional[str] = os.environ.get( + "OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1" + ) openai_embedding_api_key: Optional[str] = os.environ.get("OPENAI_EMBEDDING_API_KEY") openai_embedding_model: Optional[str] = "text-embedding-3-small" openai_chat_tokens: int = 8192 openai_extract_tokens: int = 256 openai_text2gql_tokens: int = 4096 # 2. Rerank settings - cohere_base_url: Optional[str] = os.environ.get("CO_API_URL", "https://api.cohere.com/v1/rerank") + cohere_base_url: Optional[str] = os.environ.get( + "CO_API_URL", "https://api.cohere.com/v1/rerank" + ) reranker_api_key: Optional[str] = None reranker_model: Optional[str] = None # 3. Ollama settings diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py index dfe9d1056..5fec3a778 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_config.py @@ -31,12 +31,15 @@ class BaseConfig(BaseSettings): class Config: env_file = env_path case_sensitive = False - extra = 'ignore' # ignore extra fields to avoid ValidationError + extra = "ignore" # ignore extra fields to avoid ValidationError env_ignore_empty = True def generate_env(self): if os.path.exists(env_path): - log.info("%s already exists, do you want to override with the default configuration? (y/n)", env_path) + log.info( + "%s already exists, do you want to override with the default configuration? (y/n)", + env_path, + ) update = input() if update.lower() != "y": return @@ -96,8 +99,12 @@ def _sync_env_to_object(self, env_config, config_dict): obj_value_str = str(obj_value) if obj_value is not None else "" if env_value != obj_value_str: - log.info("Update configuration from the file: %s=%s (Original value: %s)", - env_key, env_value, obj_value_str) + log.info( + "Update configuration from the file: %s=%s (Original value: %s)", + env_key, + env_value, + obj_value_str, + ) # Update the object attribute (using lowercase key) setattr(self, env_key.lower(), env_value) @@ -106,8 +113,11 @@ def _sync_object_to_env(self, env_config, config_dict): for obj_key, obj_value in config_dict.items(): if obj_key not in env_config: obj_value_str = str(obj_value) if obj_value is not None else "" - log.info("Add configuration items to the environment variable file: %s=%s", - obj_key, obj_value) + log.info( + "Add configuration items to the environment variable file: %s=%s", + obj_key, + obj_value, + ) # Add to .env set_key(env_path, obj_key, obj_value_str, quote_mode="never") diff --git a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py index 1008c3c13..4b0c4dc76 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py @@ -32,11 +32,14 @@ class LiteralStr(str): pass + def literal_str_representer(dumper, data): - return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|') + return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|") + yaml.add_representer(LiteralStr, literal_str_representer) + class BasePromptConfig: graph_schema: str = "" extract_graph_prompt: str = "" @@ -54,9 +57,7 @@ def ensure_yaml_file_exists(self): current_dir = Path.cwd().resolve() project_root = get_project_root() if current_dir == project_root: - log.info( - "Current working directory is the project root, proceeding to run the app." - ) + log.info("Current working directory is the project root, proceeding to run the app.") else: error_msg = ( f"Current working directory is not the project root. " @@ -74,16 +75,20 @@ def ensure_yaml_file_exists(self): setattr(self, key, value) # Check if the language in the .env file matches the language in the YAML file - env_lang = (self.llm_settings.language.lower() - if hasattr(self, 'llm_settings') and self.llm_settings.language - else 'en') - yaml_lang = data.get('_language_generated', 'en').lower() + env_lang = ( + self.llm_settings.language.lower() + if hasattr(self, "llm_settings") and self.llm_settings.language + else "en" + ) + yaml_lang = data.get("_language_generated", "en").lower() if env_lang.strip() != yaml_lang.strip(): log.warning( "Prompt was changed '.env' language is '%s', " "but '%s' was generated for '%s'. " "Regenerating the prompt file...", - env_lang, F_NAME, yaml_lang + env_lang, + F_NAME, + yaml_lang, ) if self.llm_settings.language.lower() == "cn": self.answer_prompt = self.answer_prompt_CN @@ -105,6 +110,7 @@ def save_to_yaml(self): def to_literal(val): return LiteralStr(val) if isinstance(val, str) else val + data = { "graph_schema": to_literal(self.graph_schema), "text2gql_graph_schema": to_literal(self.text2gql_graph_schema), diff --git a/hugegraph-llm/src/hugegraph_llm/config/prompt_config.py b/hugegraph-llm/src/hugegraph_llm/config/prompt_config.py index eaccbefa2..cc79b3cef 100644 --- a/hugegraph-llm/src/hugegraph_llm/config/prompt_config.py +++ b/hugegraph-llm/src/hugegraph_llm/config/prompt_config.py @@ -23,6 +23,7 @@ class PromptConfig(BasePromptConfig): def __init__(self, llm_config_object): self.llm_settings = llm_config_object + # Data is detached from llm_op/answer_synthesize.py answer_prompt_EN: str = """You are an expert in the fields of knowledge graphs and natural language processing. diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 2d5937a43..1b2032b23 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -30,7 +30,7 @@ async def log_stream(log_path: str, lines: int = 125): Stream the content of a log file like `tail -f`. """ try: - with open(log_path, 'r', encoding='utf-8') as file: + with open(log_path, "r", encoding="utf-8") as file: buffer = deque(file, maxlen=lines) for line in buffer: yield line # Yield the initial lines @@ -50,8 +50,8 @@ async def log_stream(log_path: str, lines: int = 125): def read_llm_server_log(lines=250): log_path = "logs/llm-server.log" try: - with open(log_path, "r", encoding='utf-8', errors="replace") as f: - return ''.join(deque(f, maxlen=lines)) + with open(log_path, "r", encoding="utf-8", errors="replace") as f: + return "".join(deque(f, maxlen=lines)) except FileNotFoundError: log.critical("Log file not found: %s", log_path) return "LLM Server log file not found." @@ -61,10 +61,10 @@ def read_llm_server_log(lines=250): def clear_llm_server_log(): log_path = "logs/llm-server.log" try: - with open(log_path, "w", encoding='utf-8') as f: + with open(log_path, "w", encoding="utf-8") as f: f.truncate(0) # Clear the contents of the file return "LLM Server log cleared." - except Exception as e: #pylint: disable=W0718 + except Exception as e: # pylint: disable=W0718 log.error("An error occurred while clearing the log: %s", str(e)) return "Failed to clear LLM Server log." @@ -84,7 +84,7 @@ def check_password(password, request: Request = None): gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), - gr.update(visible=False) + gr.update(visible=False), ) # Log the failed attempt with IP address log.error("Incorrect password attempt from IP: %s", client_ip) @@ -93,7 +93,7 @@ def check_password(password, request: Request = None): gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), - gr.update(value="Incorrect password. Access denied.", visible=True) + gr.update(value="Incorrect password. Access denied.", visible=True), ) @@ -110,10 +110,7 @@ def create_admin_block(): # Error message box, initially hidden error_message = gr.Textbox( - label="", - visible=False, - interactive=False, - elem_classes="error-message" + label="", visible=False, interactive=False, elem_classes="error-message" ) # Button to submit password @@ -136,26 +133,32 @@ def create_admin_block(): clear_llm_server_button = gr.Button("Clear LLM Server Log", visible=False) with gr.Column(): # Button to refresh LLM Server log manually - refresh_llm_server_button = gr.Button("Refresh LLM Server Log", visible=False, - variant="primary") + refresh_llm_server_button = gr.Button( + "Refresh LLM Server Log", visible=False, variant="primary" + ) # Define what happens when the password is submitted - submit_button.click( #pylint: disable=E1101 + submit_button.click( # pylint: disable=E1101 fn=check_password, inputs=[password_input], - outputs=[llm_server_log_output, hidden_row, clear_llm_server_button, - refresh_llm_server_button, error_message], + outputs=[ + llm_server_log_output, + hidden_row, + clear_llm_server_button, + refresh_llm_server_button, + error_message, + ], ) # Define what happens when the Clear LLM Server Log button is clicked - clear_llm_server_button.click( #pylint: disable=E1101 + clear_llm_server_button.click( # pylint: disable=E1101 fn=clear_llm_server_log, inputs=[], outputs=[llm_server_log_output], ) # Define what happens when the Refresh LLM Server Log button is clicked - refresh_llm_server_button.click( #pylint: disable=E1101 + refresh_llm_server_button.click( # pylint: disable=E1101 fn=read_llm_server_log, inputs=[], outputs=[llm_server_log_output], 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 4e575dddd..4e3f4de39 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -166,9 +166,7 @@ def create_app(): # settings.check_env() prompt.update_yaml_file() auth_enabled = admin_settings.enable_login.lower() == "true" - log.info( - "(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled" - ) + log.info("(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled") api_auth = APIRouter(dependencies=[Depends(authenticate)] if auth_enabled else []) hugegraph_llm = init_rag_ui() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 01ea24aa8..8c595c30d 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -71,9 +71,7 @@ def test_api_connection( log.debug("Request URL: %s", url) try: if method.upper() == "GET": - resp = requests.get( - url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth - ) + resp = requests.get(url, headers=headers, params=params, timeout=(1.0, 5.0), auth=auth) elif method.upper() == "POST": resp = requests.post( url, @@ -125,9 +123,7 @@ def apply_embedding_config(arg1, arg2, arg3, origin_call=None) -> int: llm_settings.ollama_embedding_host = arg1 llm_settings.ollama_embedding_port = int(arg2) llm_settings.ollama_embedding_model = arg3 - status_code = test_api_connection( - f"http://{arg1}:{arg2}", origin_call=origin_call - ) + status_code = test_api_connection(f"http://{arg1}:{arg2}", origin_call=origin_call) elif embedding_option == "litellm": llm_settings.litellm_embedding_api_key = arg1 llm_settings.litellm_embedding_api_base = arg2 @@ -218,8 +214,7 @@ def apply_llm_config( setattr(llm_settings, f"openai_{current_llm_config}_tokens", int(max_tokens)) test_url = ( - getattr(llm_settings, f"openai_{current_llm_config}_api_base") - + "/chat/completions" + getattr(llm_settings, f"openai_{current_llm_config}_api_base") + "/chat/completions" ) data = { "model": model_name, @@ -233,9 +228,7 @@ def apply_llm_config( elif llm_option == "ollama/local": setattr(llm_settings, f"ollama_{current_llm_config}_host", api_key_or_host) - setattr( - llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port) - ) + setattr(llm_settings, f"ollama_{current_llm_config}_port", int(api_base_or_port)) setattr(llm_settings, f"ollama_{current_llm_config}_language_model", model_name) status_code = test_api_connection( f"http://{api_key_or_host}:{api_base_or_port}", origin_call=origin_call @@ -243,12 +236,8 @@ def apply_llm_config( elif llm_option == "litellm": setattr(llm_settings, f"litellm_{current_llm_config}_api_key", api_key_or_host) - setattr( - llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port - ) - setattr( - llm_settings, f"litellm_{current_llm_config}_language_model", model_name - ) + setattr(llm_settings, f"litellm_{current_llm_config}_api_base", api_base_or_port) + setattr(llm_settings, f"litellm_{current_llm_config}_language_model", model_name) setattr(llm_settings, f"litellm_{current_llm_config}_tokens", int(max_tokens)) status_code = test_litellm_chat( @@ -295,7 +284,9 @@ def create_configs_block() -> list: ), ] graph_config_button = gr.Button("Apply Configuration") - graph_config_button.click(apply_graph_config, inputs=graph_config_input) # pylint: disable=no-member + graph_config_button.click( + apply_graph_config, inputs=graph_config_input + ) # pylint: disable=no-member # TODO : use OOP to refactor the following code with gr.Accordion("2. Set up the LLM.", open=False): @@ -373,13 +364,9 @@ def chat_llm_settings(llm_type): ), ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_chat_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_chat_op, inputs=llm_config_input) # Determine whether there are Settings in the.env file env_path = os.path.join( os.getcwd(), ".env" @@ -419,9 +406,7 @@ def extract_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr( - llm_settings, "openai_extract_language_model" - ), + value=getattr(llm_settings, "openai_extract_language_model"), label="model_name", ), gr.Textbox( @@ -440,9 +425,7 @@ def extract_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr( - llm_settings, "ollama_extract_language_model" - ), + value=getattr(llm_settings, "ollama_extract_language_model"), label="model_name", ), gr.Textbox(value="", visible=False), @@ -460,9 +443,7 @@ def extract_llm_settings(llm_type): info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr( - llm_settings, "litellm_extract_language_model" - ), + value=getattr(llm_settings, "litellm_extract_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -472,13 +453,9 @@ def extract_llm_settings(llm_type): ), ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_extract_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_extract_op, inputs=llm_config_input) with gr.Tab(label="text2gql"): text2gql_llm_dropdown = gr.Dropdown( @@ -503,9 +480,7 @@ def text2gql_llm_settings(llm_type): label="api_base", ), gr.Textbox( - value=getattr( - llm_settings, "openai_text2gql_language_model" - ), + value=getattr(llm_settings, "openai_text2gql_language_model"), label="model_name", ), gr.Textbox( @@ -524,9 +499,7 @@ def text2gql_llm_settings(llm_type): label="port", ), gr.Textbox( - value=getattr( - llm_settings, "ollama_text2gql_language_model" - ), + value=getattr(llm_settings, "ollama_text2gql_language_model"), label="model_name", ), gr.Textbox(value="", visible=False), @@ -544,9 +517,7 @@ def text2gql_llm_settings(llm_type): info="If you want to use the default api_base, please keep it blank", ), gr.Textbox( - value=getattr( - llm_settings, "litellm_text2gql_language_model" - ), + value=getattr(llm_settings, "litellm_text2gql_language_model"), label="model_name", info="Please refer to https://docs.litellm.ai/docs/providers", ), @@ -556,13 +527,9 @@ def text2gql_llm_settings(llm_type): ), ] else: - llm_config_input = [ - gr.Textbox(value="", visible=False) for _ in range(4) - ] + llm_config_input = [gr.Textbox(value="", visible=False) for _ in range(4)] llm_config_button = gr.Button("Apply configuration") - llm_config_button.click( - apply_llm_config_with_text2gql_op, inputs=llm_config_input - ) + llm_config_button.click(apply_llm_config_with_text2gql_op, inputs=llm_config_input) with gr.Accordion("3. Set up the Embedding.", open=False): embedding_dropdown = gr.Dropdown( @@ -594,12 +561,8 @@ def embedding_settings(embedding_type): elif embedding_type == "ollama/local": with gr.Row(): embedding_config_input = [ - gr.Textbox( - value=llm_settings.ollama_embedding_host, label="host" - ), - gr.Textbox( - value=str(llm_settings.ollama_embedding_port), label="port" - ), + gr.Textbox(value=llm_settings.ollama_embedding_host, label="host"), + gr.Textbox(value=str(llm_settings.ollama_embedding_port), label="port"), gr.Textbox( value=llm_settings.ollama_embedding_model, label="model_name", @@ -648,9 +611,7 @@ def embedding_settings(embedding_type): @gr.render(inputs=[reranker_dropdown]) def reranker_settings(reranker_type): - llm_settings.reranker_type = ( - reranker_type if reranker_type != "None" else None - ) + llm_settings.reranker_type = reranker_type if reranker_type != "None" else None if reranker_type == "cohere": with gr.Row(): reranker_config_input = [ @@ -660,9 +621,7 @@ def reranker_settings(reranker_type): type="password", ), gr.Textbox(value=llm_settings.reranker_model, label="model"), - gr.Textbox( - value=llm_settings.cohere_base_url, label="base_url" - ), + gr.Textbox(value=llm_settings.cohere_base_url, label="base_url"), ] elif reranker_type == "siliconflow": with gr.Row(): diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py index da10f50f4..8b78328f3 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/other_block.py @@ -31,7 +31,9 @@ def create_other_block(): gr.Markdown("""## Other Tools """) with gr.Row(): - inp = gr.Textbox(value="g.V().limit(10)", label="Gremlin query", show_copy_button=True, lines=8) + inp = gr.Textbox( + value="g.V().limit(10)", label="Gremlin query", show_copy_button=True, lines=8 + ) out = gr.Code(label="Output", language="json", elem_classes="code-container-show") btn = gr.Button("Run Gremlin query") btn.click(fn=run_gremlin_query, inputs=[inp], outputs=out) # pylint: disable=no-member @@ -39,7 +41,9 @@ def create_other_block(): gr.Markdown("---") with gr.Row(): inp = [] - out = gr.Textbox(label="Backup Graph Manually (Auto backup at 1:00 AM everyday)", show_copy_button=True) + out = gr.Textbox( + label="Backup Graph Manually (Auto backup at 1:00 AM everyday)", show_copy_button=True + ) btn = gr.Button("Backup Graph Data") btn.click(fn=backup_data, inputs=inp, outputs=out) # pylint: disable=no-member with gr.Accordion("Init HugeGraph test data (🚧)", open=False): @@ -55,10 +59,7 @@ async def lifespan(app: FastAPI): # pylint: disable=W0621 log.info("Starting background scheduler...") scheduler = AsyncIOScheduler() scheduler.add_job( - backup_data, - trigger=CronTrigger(hour=1, minute=0), - id="daily_backup", - replace_existing=True + backup_data, trigger=CronTrigger(hour=1, minute=0), id="daily_backup", replace_existing=True ) scheduler.start() diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 982436b0f..8f70c34bd 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -91,7 +91,9 @@ def rag_answer( near_neighbor_first=near_neighbor_first, topk_return_results=topk_return_results, ) - rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) + rag.synthesize_answer( + raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt + ) try: context = rag.run( @@ -146,6 +148,7 @@ def update_ui_configs( graph_search = graph_only_answer or graph_vector_answer return graph_search, gremlin_prompt, vector_search + async def rag_answer_streaming( text: str, raw_answer: bool, @@ -188,9 +191,9 @@ async def rag_answer_streaming( if vector_search: rag.query_vector_index() if graph_search: - rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid().import_schema( - huge_settings.graph_name - ).query_graphdb( + rag.extract_keywords( + extract_template=keywords_extract_prompt + ).keywords_to_vid().import_schema(huge_settings.graph_name).query_graphdb( num_gremlin_generate_example=gremlin_tmpl_num, gremlin_prompt=gremlin_prompt, ) @@ -202,7 +205,9 @@ async def rag_answer_streaming( # rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) try: - context = rag.run(verbose=True, query=text, vector_search=vector_search, graph_search=graph_search) + context = rag.run( + verbose=True, query=text, vector_search=vector_search, graph_search=graph_search + ) if context.get("switch_to_bleu"): gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") answer_synthesize = AnswerSynthesize( @@ -228,6 +233,7 @@ async def rag_answer_streaming( log.critical(e) raise gr.Error(f"An unexpected error occurred: {str(e)}") + @with_task_id def create_rag_block(): # pylint: disable=R0915 (too-many-statements),C0301 @@ -235,7 +241,9 @@ def create_rag_block(): with gr.Row(): with gr.Column(scale=2): # with gr.Blocks().queue(max_size=20, default_concurrency_limit=5): - inp = gr.Textbox(value=prompt.default_question, label="Question", show_copy_button=True, lines=3) + inp = gr.Textbox( + value=prompt.default_question, label="Question", show_copy_button=True, lines=3 + ) # TODO: Only support inline formula now. Should support block formula gr.Markdown("Basic LLM Answer", elem_classes="output-box-label") @@ -275,10 +283,16 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") - vector_only_radio = gr.Radio(choices=[True, False], value=False, label="Vector-only Answer") + vector_only_radio = gr.Radio( + choices=[True, False], value=False, label="Vector-only Answer" + ) with gr.Row(): - graph_only_radio = gr.Radio(choices=[True, False], value=True, label="Graph-only Answer") - graph_vector_radio = gr.Radio(choices=[True, False], value=False, label="Graph-Vector Answer") + graph_only_radio = gr.Radio( + choices=[True, False], value=True, label="Graph-only Answer" + ) + graph_vector_radio = gr.Radio( + choices=[True, False], value=False, label="Graph-Vector Answer" + ) def toggle_slider(enable): return gr.update(interactive=enable) @@ -291,8 +305,12 @@ def toggle_slider(enable): value="reranker" if online_rerank else "bleu", label="Rerank method", ) - example_num = gr.Number(value=-1, label="Template Num (<0 means disable text2gql) ", precision=0) - graph_ratio = gr.Slider(0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False) + example_num = gr.Number( + value=-1, label="Template Num (<0 means disable text2gql) ", precision=0 + ) + graph_ratio = gr.Slider( + 0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False + ) graph_vector_radio.change( toggle_slider, inputs=graph_vector_radio, outputs=graph_ratio @@ -325,8 +343,8 @@ def toggle_slider(enable): example_num, ], outputs=[raw_out, vector_only_out, graph_only_out, graph_vector_out], - queue=True, # Enable queueing for this event - concurrency_limit=5, # Maximum of 5 concurrent executions + queue=True, # Enable queueing for this event + concurrency_limit=5, # Maximum of 5 concurrent executions ) gr.Markdown( @@ -394,18 +412,20 @@ def several_rag_answer( total_rows = len(df) for index, row in df.iterrows(): question = row.iloc[0] - basic_llm_answer, vector_only_answer, graph_only_answer, graph_vector_answer = rag_answer( - question, - is_raw_answer, - is_vector_only_answer, - is_graph_only_answer, - is_graph_vector_answer, - graph_ratio_ui, - rerank_method_ui, - near_neighbor_first_ui, - custom_related_information_ui, - answer_prompt, - keywords_extract_prompt, + basic_llm_answer, vector_only_answer, graph_only_answer, graph_vector_answer = ( + rag_answer( + question, + is_raw_answer, + is_vector_only_answer, + is_graph_only_answer, + is_graph_vector_answer, + graph_ratio_ui, + rerank_method_ui, + near_neighbor_first_ui, + custom_related_information_ui, + answer_prompt, + keywords_extract_prompt, + ) ) df.at[index, "Basic LLM Answer"] = basic_llm_answer df.at[index, "Vector-only Answer"] = vector_only_answer @@ -418,7 +438,9 @@ def several_rag_answer( with gr.Row(): with gr.Column(): - questions_file = gr.File(file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)") + questions_file = gr.File( + file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" + ) with gr.Column(): test_template_file = os.path.join(resource_path, "demo", "questions_template.xlsx") gr.File(value=test_template_file, label="Download Template File") diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 7d682403f..6600d7c41 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -33,11 +33,13 @@ from hugegraph_llm.utils.embedding_utils import get_index_folder_name from hugegraph_llm.utils.hugegraph_utils import run_gremlin_query from hugegraph_llm.utils.log import log +from hugegraph_llm.flows.scheduler import SchedulerSingleton @dataclass class GremlinResult: """Standardized result class for gremlin_generate function""" + success: bool match_result: str template_gremlin: Optional[str] = None @@ -47,13 +49,19 @@ class GremlinResult: error_message: Optional[str] = None @classmethod - def error(cls, message: str) -> 'GremlinResult': + def error(cls, message: str) -> "GremlinResult": """Create an error result""" return cls(success=False, match_result=message, error_message=message) @classmethod - def success_result(cls, match_result: str, template_gremlin: str, - raw_gremlin: str, template_exec: str, raw_exec: str) -> 'GremlinResult': + def success_result( + cls, + match_result: str, + template_gremlin: str, + raw_gremlin: str, + template_exec: str, + raw_exec: str, + ) -> "GremlinResult": """Create a successful result""" return cls( success=True, @@ -61,7 +69,7 @@ def success_result(cls, match_result: str, template_gremlin: str, template_gremlin=template_gremlin, raw_gremlin=raw_gremlin, template_exec_result=template_exec, - raw_exec_result=raw_exec + raw_exec_result=raw_exec, ) @@ -93,6 +101,7 @@ def build_example_vector_index(temp_file) -> dict: target_file = os.path.join(resource_path, folder_name, "gremlin_examples", file_name) try: import shutil + shutil.copy2(full_path, target_file) log.info("Successfully copied file to: %s", target_file) except (OSError, IOError) as e: @@ -143,7 +152,7 @@ def _configure_output_types(requested_outputs): "template_gremlin": True, "raw_gremlin": True, "template_execution_result": True, - "raw_execution_result": True + "raw_execution_result": True, } if requested_outputs: for key in output_types: @@ -176,7 +185,9 @@ def _execute_queries(context, output_types): def gremlin_generate( inp, example_num, schema, gremlin_prompt, requested_outputs: Optional[List[str]] = None ) -> GremlinResult: - generator = GremlinGenerator(llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding()) + generator = GremlinGenerator( + llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() + ) sm = SchemaManager(graph_name=schema) processed_schema, short_schema = _process_schema(schema, generator, sm) @@ -196,7 +207,9 @@ def gremlin_generate( _execute_queries(context, output_types) - match_result = json.dumps(context.get("match_result", "No Results"), ensure_ascii=False, indent=2) + match_result = json.dumps( + context.get("match_result", "No Results"), ensure_ascii=False, indent=2 + ) return GremlinResult.success_result( match_result=match_result, template_gremlin=context["result"], @@ -220,7 +233,11 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: if "edgelabels" in schema: mini_schema["edgelabels"] = [] for edge in schema["edgelabels"]: - new_edge = {key: edge[key] for key in ["name", "source_label", "target_label", "properties"] if key in edge} + new_edge = { + key: edge[key] + for key in ["name", "source_label", "target_label", "properties"] + if key in edge + } mini_schema["edgelabels"].append(new_edge) return mini_schema @@ -228,17 +245,40 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: def gremlin_generate_for_ui(inp, example_num, schema, gremlin_prompt): """UI wrapper for gremlin_generate that returns tuple for Gradio compatibility""" - result = gremlin_generate(inp, example_num, schema, gremlin_prompt) - - if not result.success: - return result.match_result, "", "", "", "" + # Execute via scheduler + try: + res = SchedulerSingleton.get_instance().schedule_flow( + "text2gremlin", + inp, + int(example_num) if isinstance(example_num, (int, float, str)) else 2, + schema, + gremlin_prompt, + [ + "match_result", + "template_gremlin", + "raw_gremlin", + "template_execution_result", + "raw_execution_result", + ], + ) + except Exception as e: # pylint: disable=broad-except + log.error("UI text2gremlin error: %s", e) + return json.dumps({"error": str(e)}, ensure_ascii=False), "", "", "", "" + + # Backward-compatible mapping for outputs + match_result = res.get("match_result", []) + match_result_str = ( + json.dumps(match_result, ensure_ascii=False, indent=2) + if isinstance(match_result, (list, dict)) + else str(match_result) + ) return ( - result.match_result, - result.template_gremlin or "", - result.raw_gremlin or "", - result.template_exec_result or "", - result.raw_exec_result or "" + match_result_str, + res.get("template_gremlin", "") or "", + res.get("raw_gremlin", "") or "", + res.get("template_execution_result", "") or "", + res.get("raw_execution_result", "") or "", ) @@ -253,7 +293,8 @@ def create_text2gremlin_block() -> Tuple: ) with gr.Row(): file = gr.File( - value=os.path.join(resource_path, "demo", "text2gremlin.csv"), label="Upload Text-Gremlin Pairs File" + value=os.path.join(resource_path, "demo", "text2gremlin.csv"), + label="Upload Text-Gremlin Pairs File", ) out = gr.Textbox(label="Result Message") with gr.Row(): @@ -263,22 +304,39 @@ def create_text2gremlin_block() -> Tuple: with gr.Row(): with gr.Column(scale=1): - input_box = gr.Textbox(value=prompt.default_question, label="Nature Language Query", show_copy_button=True) - match = gr.Code(label="Similar Template (TopN)", language="javascript", elem_classes="code-container-show") + input_box = gr.Textbox( + value=prompt.default_question, label="Nature Language Query", show_copy_button=True + ) + match = gr.Code( + label="Similar Template (TopN)", + language="javascript", + elem_classes="code-container-show", + ) initialized_out = gr.Textbox(label="Gremlin With Template", show_copy_button=True) raw_out = gr.Textbox(label="Gremlin Without Template", show_copy_button=True) tmpl_exec_out = gr.Code( - label="Query With Template Output", language="json", elem_classes="code-container-show" + label="Query With Template Output", + language="json", + elem_classes="code-container-show", ) raw_exec_out = gr.Code( - label="Query Without Template Output", language="json", elem_classes="code-container-show" + label="Query Without Template Output", + language="json", + elem_classes="code-container-show", ) with gr.Column(scale=1): - example_num_slider = gr.Slider(minimum=0, maximum=10, step=1, value=2, label="Number of refer examples") - schema_box = gr.Textbox(value=prompt.text2gql_graph_schema, label="Schema", lines=2, show_copy_button=True) + example_num_slider = gr.Slider( + minimum=0, maximum=10, step=1, value=2, label="Number of refer examples" + ) + schema_box = gr.Textbox( + value=prompt.text2gql_graph_schema, label="Schema", lines=2, show_copy_button=True + ) prompt_box = gr.Textbox( - value=prompt.gremlin_generate_prompt, label="Prompt", lines=20, show_copy_button=True + value=prompt.gremlin_generate_prompt, + label="Prompt", + lines=20, + show_copy_button=True, ) btn = gr.Button("Text2Gremlin", variant="primary") btn.click( # pylint: disable=no-member @@ -324,6 +382,7 @@ def graph_rag_recall( context = rag.run(verbose=True, query=query, graph_search=True) return context + def gremlin_generate_selective( inp: str, example_num: int, diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 4aa476942..56b5de4b3 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -63,16 +63,12 @@ def generate_prompt_for_ui(source_text, scenario, example_name): Handles the UI logic for generating a new prompt using the new workflow architecture. """ if not all([source_text, scenario, example_name]): - gr.Warning( - "Please provide original text, expected scenario, and select an example!" - ) + gr.Warning("Please provide original text, expected scenario, and select an example!") return gr.update() try: # using new architecture scheduler = SchedulerSingleton.get_instance() - result = scheduler.schedule_flow( - "prompt_generate", source_text, scenario, example_name - ) + result = scheduler.schedule_flow("prompt_generate", source_text, scenario, example_name) gr.Info("Prompt generated successfully!") return result except Exception as e: @@ -83,9 +79,7 @@ def generate_prompt_for_ui(source_text, scenario, example_name): def load_example_names(): """Load all candidate examples""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "prompt_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return [example.get("name", "Unnamed example") for example in examples] @@ -99,27 +93,23 @@ def load_query_examples(): language = getattr( prompt, "language", - getattr(prompt.llm_settings, "language", "EN") - if hasattr(prompt, "llm_settings") - else "EN", + ( + getattr(prompt.llm_settings, "language", "EN") + if hasattr(prompt, "llm_settings") + else "EN" + ), ) if language.upper() == "CN": - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples_CN.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples_CN.json") else: - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) except (FileNotFoundError, json.JSONDecodeError): try: - examples_path = os.path.join( - resource_path, "prompt_examples", "query_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -130,9 +120,7 @@ def load_query_examples(): def load_schema_fewshot_examples(): """Load few-shot examples from a JSON file""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "schema_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "schema_examples.json") with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -143,14 +131,10 @@ def load_schema_fewshot_examples(): def update_example_preview(example_name): """Update the display content based on the selected example name.""" try: - examples_path = os.path.join( - resource_path, "prompt_examples", "prompt_examples.json" - ) + examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") with open(examples_path, "r", encoding="utf-8") as f: all_examples = json.load(f) - selected_example = next( - (ex for ex in all_examples if ex.get("name") == example_name), None - ) + selected_example = next((ex for ex in all_examples if ex.get("name") == example_name), None) if selected_example: return ( @@ -178,9 +162,11 @@ def _create_prompt_helper_block(demo, input_text, info_extract_template): few_shot_dropdown = gr.Dropdown( choices=example_names, label="Select a Few-shot example as a reference", - value=example_names[0] - if example_names and example_names[0] != "No available examples" - else None, + value=( + example_names[0] + if example_names and example_names[0] != "No available examples" + else None + ), ) with gr.Accordion("View example details", open=False): example_desc_preview = gr.Markdown(label="Example description") @@ -193,9 +179,7 @@ def _create_prompt_helper_block(demo, input_text, info_extract_template): interactive=False, ) - generate_prompt_btn = gr.Button( - "🚀 Auto-generate Graph Extract Prompt", variant="primary" - ) + generate_prompt_btn = gr.Button("🚀 Auto-generate Graph Extract Prompt", variant="primary") # Bind the change event of the dropdown menu few_shot_dropdown.change( fn=update_example_preview, @@ -287,9 +271,7 @@ def create_vector_graph_block(): lines=15, max_lines=29, ) - out = gr.Code( - label="Output Info", language="json", elem_classes="code-container-edit" - ) + out = gr.Code(label="Output Info", language="json", elem_classes="code-container-edit") with gr.Row(): with gr.Accordion("Get RAG Info", open=False): @@ -298,12 +280,8 @@ def create_vector_graph_block(): graph_index_btn0 = gr.Button("Get Graph Index Info", size="sm") with gr.Accordion("Clear RAG Data", open=False): with gr.Column(): - vector_index_btn1 = gr.Button( - "Clear Chunks Vector Index", size="sm" - ) - graph_index_btn1 = gr.Button( - "Clear Graph Vid Vector Index", size="sm" - ) + vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") + graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") vector_import_bt = gr.Button("Import into Vector", variant="primary") @@ -376,9 +354,9 @@ def create_vector_graph_block(): inputs=[input_text, input_schema, info_extract_template], ) - graph_loading_bt.click( - import_graph_data, inputs=[out, input_schema], outputs=[out] - ).then(update_vid_embedding).then( + graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( + update_vid_embedding + ).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) diff --git a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py index 495ef667c..ee173b284 100644 --- a/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py +++ b/hugegraph-llm/src/hugegraph_llm/document/chunk_split.py @@ -22,9 +22,9 @@ class ChunkSplitter: def __init__( - self, - split_type: Literal["paragraph", "sentence"] = "paragraph", - language: Literal["zh", "en"] = "zh" + self, + split_type: Literal["paragraph", "sentence"] = "paragraph", + language: Literal["zh", "en"] = "zh", ): if language == "zh": separators = ["\n\n", "\n", "。", "，", ""] @@ -34,15 +34,11 @@ def __init__( raise ValueError("Argument `language` must be zh or en!") if split_type == "paragraph": self.text_splitter = RecursiveCharacterTextSplitter( - chunk_size=500, - chunk_overlap=30, - separators=separators + chunk_size=500, chunk_overlap=30, separators=separators ) elif split_type == "sentence": self.text_splitter = RecursiveCharacterTextSplitter( - chunk_size=50, - chunk_overlap=0, - separators=separators + chunk_size=50, chunk_overlap=0, separators=separators ) else: raise ValueError("Arg `type` must be paragraph, sentence!") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py index fa10d0199..7d2735352 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py @@ -48,9 +48,7 @@ def build_flow(self, *args, **kwargs): def post_deal(self, pipeline=None): graph_summary_info = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) filename_prefix = get_filename_prefix( llm_settings.embedding_type, diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py index 1b0c98253..55f53b7ad 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py @@ -27,9 +27,7 @@ class GraphExtractFlow(BaseFlow): def __init__(self): pass - def prepare( - self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type - ): + def prepare(self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type): # prepare input data prepared_input.texts = texts prepared_input.language = "zh" 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 5581ef107..0b29b4e64 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py @@ -35,9 +35,11 @@ def prepare(self, prepared_input: WkFlowInput, data, schema): raise ValueError(f"Invalid JSON for 'data': {e.msg}") from e log.debug( "Import graph data (truncated): %s", - (data[:512] + "...") - if isinstance(data, str) and len(data) > 512 - else (data if isinstance(data, str) else ""), + ( + (data[:512] + "...") + if isinstance(data, str) and len(data) > 512 + else (data if isinstance(data, str) else "") + ), ) prepared_input.data_json = data_json prepared_input.schema = schema diff --git a/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py index aece6bd61..b4a7bf329 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py @@ -58,6 +58,4 @@ def post_deal(self, pipeline=None): Process the execution result of PromptGenerate workflow """ res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - return res.get( - "generated_extract_prompt", "Generation failed. Please check the logs." - ) + return res.get("generated_extract_prompt", "Generation failed. Please check the logs.") diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index 559540ce3..3aedbe7f2 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -25,6 +25,7 @@ from hugegraph_llm.flows.build_schema import BuildSchemaFlow from hugegraph_llm.flows.prompt_generate import PromptGenerateFlow from hugegraph_llm.utils.log import log +from hugegraph_llm.flows.text2gremlin import Text2GremlinFlow class Scheduler: @@ -62,6 +63,10 @@ def __init__(self, max_pipeline: int = 10): "manager": GPipelineManager(), "flow": PromptGenerateFlow(), } + self.pipeline_pool["text2gremlin"] = { + "manager": GPipelineManager(), + "flow": Text2GremlinFlow(), + } self.max_pipeline = max_pipeline # TODO: Implement Agentic Workflow @@ -96,7 +101,9 @@ def schedule_flow(self, flow: str, *args, **kwargs): flow.prepare(prepared_input, *args, **kwargs) status = pipeline.run() if status.isErr(): - raise RuntimeError(f"Error in flow execution {status.getInfo()}") + error_msg = f"Error in flow execution {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) res = flow.post_deal(pipeline) manager.release(pipeline) return res diff --git a/hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py new file mode 100644 index 000000000..e9ba4276c --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py @@ -0,0 +1,112 @@ +# 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 PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.nodes.index_node.gremlin_example_index_query import GremlinExampleIndexQueryNode +from hugegraph_llm.nodes.llm_node.text2gremlin import Text2GremlinNode +from hugegraph_llm.nodes.hugegraph_node.gremlin_execute import GremlinExecuteNode + +from typing import Any, Dict, List, Optional + + +class Text2GremlinFlow(BaseFlow): + def __init__(self): + pass + + def prepare( + self, + prepared_input: WkFlowInput, + query: str, + example_num: int, + schema_input: str, + gremlin_prompt_input: Optional[str], + requested_outputs: Optional[List[str]], + ): + # sanitize example_num to [0,10], fallback to 2 if invalid + if not isinstance(example_num, int): + example_num = 2 + example_num = max(0, min(10, example_num)) + + # filter requested_outputs to allowed set and cap to 5 + allowed = { + "match_result", + "template_gremlin", + "raw_gremlin", + "template_execution_result", + "raw_execution_result", + } + req = requested_outputs or ["template_gremlin"] + req = [x for x in req if x in allowed] + if not req: + req = ["template_gremlin"] + if len(req) > 5: + req = req[:5] + + prepared_input.query = query + prepared_input.example_num = example_num + prepared_input.schema = schema_input + prepared_input.gremlin_prompt = gremlin_prompt_input + prepared_input.requested_outputs = req + return + + def build_flow( + self, + query: str, + example_num: int, + schema_input: str, + gremlin_prompt_input: Optional[str] = None, + requested_outputs: Optional[List[str]] = None, + ): + pipeline = GPipeline() + + prepared_input = WkFlowInput() + self.prepare( + prepared_input, + query=query, + example_num=example_num, + schema_input=schema_input, + gremlin_prompt_input=gremlin_prompt_input, + requested_outputs=requested_outputs, + ) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + schema_node = SchemaNode() + ieq_node = GremlinExampleIndexQueryNode() + tgn_node = Text2GremlinNode() + exe_node = GremlinExecuteNode() + + pipeline.registerGElement(schema_node, set(), "schema_node") + pipeline.registerGElement(ieq_node, set(), "gremlin_example_index_query") + pipeline.registerGElement(tgn_node, {schema_node, ieq_node}, "text2gremlin") + pipeline.registerGElement(exe_node, {tgn_node}, "gremlin_execute") + + return pipeline + + def post_deal(self, pipeline=None) -> Dict[str, Any]: + state = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + # 始终返回 5 个标准键，避免前端因过滤异常看不到字段 + return { + "match_result": state.get("match_result", []), + "template_gremlin": state.get("result", ""), + "raw_gremlin": state.get("raw_result", ""), + "template_execution_result": state.get("template_exec_res", ""), + "raw_execution_result": state.get("raw_exec_res", ""), + } diff --git a/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py b/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py index e78aa6d58..694ca014d 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/graph_index.py @@ -24,15 +24,16 @@ class GraphIndex: def __init__( - self, - graph_url: Optional[str] = huge_settings.graph_url, - graph_name: Optional[str] = huge_settings.graph_name, - graph_user: Optional[str] = huge_settings.graph_user, - graph_pwd: Optional[str] = huge_settings.graph_pwd, - graph_space: Optional[str] = huge_settings.graph_space, + self, + graph_url: Optional[str] = huge_settings.graph_url, + graph_name: Optional[str] = huge_settings.graph_name, + graph_user: Optional[str] = huge_settings.graph_user, + graph_pwd: Optional[str] = huge_settings.graph_pwd, + graph_space: Optional[str] = huge_settings.graph_space, ): - self.client = PyHugeClient(url=graph_url, graph=graph_name, user=graph_user, pwd=graph_pwd, - graphspace=graph_space) + self.client = PyHugeClient( + url=graph_url, graph=graph_name, user=graph_user, pwd=graph_pwd, graphspace=graph_space + ) def clear_graph(self): self.client.gremlin().exec("g.V().drop()") diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py index 641ac6d6e..f85483185 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index.py @@ -37,7 +37,9 @@ def __init__(self, embed_dim: int = 1024): self.properties = [] @staticmethod - def from_index_file(dir_path: str, filename_prefix: str = None, record_miss: bool = True) -> "VectorIndex": + def from_index_file( + dir_path: str, filename_prefix: str = None, record_miss: bool = True + ) -> "VectorIndex": """Load index from files, supporting model-specific filenames. This method loads a Faiss index and its corresponding properties from a directory. @@ -47,13 +49,18 @@ def from_index_file(dir_path: str, filename_prefix: str = None, record_miss: boo matches the number of properties. """ index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + property_name = ( + f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + ) index_file = os.path.join(dir_path, index_name) properties_file = os.path.join(dir_path, property_name) miss_files = [f for f in [index_file, properties_file] if not os.path.exists(f)] if miss_files: if record_miss: - log.warning("Missing vector files: %s. \nNeed create a new one for it.", ", ".join(miss_files)) + log.warning( + "Missing vector files: %s. \nNeed create a new one for it.", + ", ".join(miss_files), + ) return VectorIndex() try: @@ -61,7 +68,9 @@ def from_index_file(dir_path: str, filename_prefix: str = None, record_miss: boo with open(properties_file, "rb") as f: properties = pkl.load(f) except (RuntimeError, pkl.UnpicklingError, OSError) as e: - log.error("Failed to load index files for model '%s': %s", filename_prefix or "default", e) + log.error( + "Failed to load index files for model '%s': %s", filename_prefix or "default", e + ) raise RuntimeError( f"Could not load index files for model '{filename_prefix or 'default'}'. " f"Original error ({type(e).__name__}): {e}" @@ -85,7 +94,9 @@ def to_index_file(self, dir_path: str, filename_prefix: str = None): os.makedirs(dir_path) index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + property_name = ( + f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + ) index_file = os.path.join(dir_path, index_name) properties_file = os.path.join(dir_path, property_name) faiss.write_index(self.index, index_file) @@ -115,7 +126,9 @@ def remove(self, props: Union[Set[Any], List[Any]]) -> int: self.properties = [p for i, p in enumerate(self.properties) if i not in indices] return remove_num - def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9) -> List[Any]: + def search( + self, query_vector: List[float], top_k: int, dis_threshold: float = 0.9 + ) -> List[Any]: if self.index.ntotal == 0: return [] @@ -129,7 +142,9 @@ def search(self, query_vector: List[float], top_k: int, dis_threshold: float = 0 results.append(deepcopy(self.properties[i])) log.debug("[✓] Add valid distance %s to results.", dist) else: - log.debug("[x] Distance %s >= threshold %s, ignore this result.", dist, dis_threshold) + log.debug( + "[x] Distance %s >= threshold %s, ignore this result.", dist, dis_threshold + ) return results @staticmethod @@ -140,7 +155,9 @@ def clean(dir_path: str, filename_prefix: str = None): If model_name is None, it targets the default files. """ index_name = f"{filename_prefix}_{INDEX_FILE_NAME}" if filename_prefix else INDEX_FILE_NAME - property_name = f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + property_name = ( + f"{filename_prefix}_{PROPERTIES_FILE_NAME}" if filename_prefix else PROPERTIES_FILE_NAME + ) index_file = os.path.join(dir_path, index_name) properties_file = os.path.join(dir_path, property_name) diff --git a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py index 47c70e1a4..c73242012 100644 --- a/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py +++ b/hugegraph-llm/src/hugegraph_llm/middleware/middleware.py @@ -26,6 +26,7 @@ # TODO: we could use middleware(AOP) in the future (dig out the lifecycle of gradio & fastapi) class UseTimeMiddleware(BaseHTTPMiddleware): """Middleware to add process time to response headers""" + def __init__(self, app): super().__init__(app) @@ -33,7 +34,7 @@ async def dispatch(self, request: Request, call_next): # TODO: handle time record for async task pool in gradio start_time = time.perf_counter() response = await call_next(request) - process_time = (time.perf_counter() - start_time) * 1000 # ms + process_time = (time.perf_counter() - start_time) * 1000 # ms unit = "ms" if process_time > 1000: process_time /= 1000 @@ -46,6 +47,6 @@ async def dispatch(self, request: Request, call_next): request.method, request.query_params, request.client.host, - request.url + request.url, ) return response diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py index db9b2f105..698b92837 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py @@ -32,9 +32,9 @@ class SimilarityMode(str, Enum): def similarity( - embedding1: Union[List[float], np.ndarray], - embedding2: Union[List[float], np.ndarray], - mode: SimilarityMode = SimilarityMode.DEFAULT, + embedding1: Union[List[float], np.ndarray], + embedding2: Union[List[float], np.ndarray], + mode: SimilarityMode = SimilarityMode.DEFAULT, ) -> float: """Get embedding similarity.""" if isinstance(embedding1, list): @@ -57,28 +57,22 @@ class BaseEmbedding(ABC): # TODO: replace all the usage by get_texts_embeddings() & remove it in the future @deprecated("Use get_texts_embeddings() instead in the future.") @abstractmethod - def get_text_embedding( - self, - text: str - ) -> List[float]: + def get_text_embedding(self, text: str) -> List[float]: """Comment""" @abstractmethod - def get_texts_embeddings( - self, - texts: List[str] - ) -> List[List[float]]: + def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch. - + This method should efficiently process multiple texts at once by leveraging the embedding model's batching capabilities, which is typically more efficient than processing texts individually. - + Parameters ---------- texts : List[str] A list of text strings to be embedded. - + Returns ------- List[List[float]] @@ -87,12 +81,9 @@ def get_texts_embeddings( """ @abstractmethod - async def async_get_texts_embeddings( - self, - texts: List[str] - ) -> List[List[float]]: + async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch asynchronously. - + This method should efficiently process multiple texts at once by leveraging the embedding model's batching capabilities, which is typically more efficient than processing texts individually. @@ -101,7 +92,7 @@ async def async_get_texts_embeddings( ---------- texts : List[str] A list of text strings to be embedded. - + Returns ------- List[List[float]] @@ -111,9 +102,9 @@ async def async_get_texts_embeddings( @staticmethod def similarity( - embedding1: Union[List[float], np.ndarray], - embedding2: Union[List[float], np.ndarray], - mode: SimilarityMode = SimilarityMode.DEFAULT, + embedding1: Union[List[float], np.ndarray], + embedding2: Union[List[float], np.ndarray], + mode: SimilarityMode = SimilarityMode.DEFAULT, ) -> float: """Get embedding similarity.""" if isinstance(embedding1, list): diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index f4026ad7f..d0e15f000 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -23,12 +23,12 @@ class OpenAIEmbedding: def __init__( - self, - model_name: str = "text-embedding-3-small", - api_key: Optional[str] = None, - api_base: Optional[str] = None + self, + model_name: str = "text-embedding-3-small", + api_key: Optional[str] = None, + api_base: Optional[str] = None, ): - api_key = api_key or '' + api_key = api_key or "" self.client = OpenAI(api_key=api_key, base_url=api_base) self.aclient = AsyncOpenAI(api_key=api_key, base_url=api_base) self.model_name = model_name @@ -38,21 +38,18 @@ def get_text_embedding(self, text: str) -> List[float]: response = self.client.embeddings.create(input=text, model=self.model_name) return response.data[0].embedding - def get_texts_embeddings( - self, - texts: List[str] - ) -> List[List[float]]: + def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch. - + This method efficiently processes multiple texts at once by leveraging OpenAI's batching capabilities, which is more efficient than processing texts individually. - + Parameters ---------- texts : List[str] A list of text strings to be embedded. - + Returns ------- List[List[float]] @@ -64,7 +61,7 @@ def get_texts_embeddings( async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """Get embeddings for multiple texts in a single batch asynchronously. - + This method should efficiently process multiple texts at once by leveraging the embedding model's batching capabilities, which is typically more efficient than processing texts individually. @@ -73,7 +70,7 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] ---------- texts : List[str] A list of text strings to be embedded. - + Returns ------- List[List[float]] diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/base.py b/hugegraph-llm/src/hugegraph_llm/models/llms/base.py index c6bfa44a8..69c082690 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/base.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/base.py @@ -24,48 +24,48 @@ class BaseLLM(ABC): @abstractmethod def generate( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, ) -> str: """Comment""" @abstractmethod async def agenerate( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, ) -> str: """Comment""" @abstractmethod def generate_streaming( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, - on_token_callback: Optional[Callable] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, + on_token_callback: Optional[Callable] = None, ) -> Generator[str, None, None]: """Comment""" @abstractmethod async def agenerate_streaming( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, - on_token_callback: Optional[Callable] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, + on_token_callback: Optional[Callable] = None, ) -> AsyncGenerator[str, None]: """Comment""" @abstractmethod def num_tokens_from_string( - self, - string: str, + self, + string: str, ) -> str: """Given a string returns the number of tokens the given string consists of""" @abstractmethod def max_allowed_token_length( - self, + self, ) -> int: """Returns the maximum number of tokens the LLM can handle""" 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 7e1eaab68..9121fca09 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py @@ -173,8 +173,4 @@ def get_text2gql_llm(self): if __name__ == "__main__": client = LLMs().get_chat_llm() print(client.generate(prompt="What is the capital of China?")) - print( - client.generate( - messages=[{"role": "user", "content": "What is the capital of China?"}] - ) - ) + print(client.generate(messages=[{"role": "user", "content": "What is the capital of China?"}])) diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/litellm.py b/hugegraph-llm/src/hugegraph_llm/models/llms/litellm.py index b9cc0f19f..6f3c8129c 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/litellm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/litellm.py @@ -51,7 +51,7 @@ def __init__( @retry( stop=stop_after_attempt(2), wait=wait_exponential(multiplier=1, min=2, max=5), - retry=retry_if_exception_type((RateLimitError, BudgetExceededError, APIError)) + retry=retry_if_exception_type((RateLimitError, BudgetExceededError, APIError)), ) def generate( self, @@ -80,12 +80,12 @@ def generate( @retry( stop=stop_after_attempt(2), wait=wait_exponential(multiplier=1, min=2, max=5), - retry=retry_if_exception_type((RateLimitError, BudgetExceededError, APIError)) + retry=retry_if_exception_type((RateLimitError, BudgetExceededError, APIError)), ) async def agenerate( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, ) -> str: """Generate a response to the query messages/prompt asynchronously.""" if messages is None: diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py index 5354ba306..6d08ce8cd 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/ollama.py @@ -28,6 +28,7 @@ class OllamaClient(BaseLLM): """LLM wrapper should take in a prompt and return a string.""" + def __init__(self, model: str, host: str = "127.0.0.1", port: int = 11434, **kwargs): self.model = model self.client = ollama.Client(host=f"http://{host}:{port}", **kwargs) @@ -49,9 +50,9 @@ def generate( messages=messages, ) usage = { - "prompt_tokens": response['prompt_eval_count'], - "completion_tokens": response['eval_count'], - "total_tokens": response['prompt_eval_count'] + response['eval_count'], + "prompt_tokens": response["prompt_eval_count"], + "completion_tokens": response["eval_count"], + "total_tokens": response["prompt_eval_count"] + response["eval_count"], } log.info("Token usage: %s", json.dumps(usage)) return response["message"]["content"] @@ -61,9 +62,9 @@ def generate( @retry(tries=3, delay=1) async def agenerate( - self, - messages: Optional[List[Dict[str, Any]]] = None, - prompt: Optional[str] = None, + self, + messages: Optional[List[Dict[str, Any]]] = None, + prompt: Optional[str] = None, ) -> str: """Comment""" if messages is None: @@ -75,9 +76,9 @@ async def agenerate( messages=messages, ) usage = { - "prompt_tokens": response['prompt_eval_count'], - "completion_tokens": response['eval_count'], - "total_tokens": response['prompt_eval_count'] + response['eval_count'], + "prompt_tokens": response["prompt_eval_count"], + "completion_tokens": response["eval_count"], + "total_tokens": response["prompt_eval_count"] + response["eval_count"], } log.info("Token usage: %s", json.dumps(usage)) return response["message"]["content"] @@ -96,11 +97,7 @@ def generate_streaming( assert prompt is not None, "Messages or prompt must be provided." messages = [{"role": "user", "content": prompt}] - for chunk in self.client.chat( - model=self.model, - messages=messages, - stream=True - ): + for chunk in self.client.chat(model=self.model, messages=messages, stream=True): if not chunk["message"]: log.debug("Received empty chunk['message'] in streaming chunk: %s", chunk) continue @@ -122,9 +119,7 @@ async def agenerate_streaming( try: async_generator = await self.async_client.chat( - model=self.model, - messages=messages, - stream=True + model=self.model, messages=messages, stream=True ) async for chunk in async_generator: token = chunk.get("message", {}).get("content", "") @@ -135,7 +130,6 @@ async def agenerate_streaming( print(f"Retrying LLM call {e}") raise e - def num_tokens_from_string( self, string: str, diff --git a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py index 88cea3976..e1088c890 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/openai.py @@ -42,7 +42,7 @@ def __init__( max_tokens: int = 8092, temperature: float = 0.01, ) -> None: - api_key = api_key or '' + api_key = api_key or "" self.client = OpenAI(api_key=api_key, base_url=api_base) self.aclient = AsyncOpenAI(api_key=api_key, base_url=api_base) self.model = model_name @@ -186,7 +186,7 @@ async def agenerate_streaming( temperature=self.temperature, max_tokens=self.max_tokens, messages=messages, - stream=True + stream=True, ) async for chunk in completions: if not chunk.choices: diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py index 1710acfc2..3bf481ce2 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/cohere.py @@ -31,16 +31,21 @@ def __init__( self.base_url = base_url self.model = model - def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int] = None) -> List[str]: + def get_rerank_lists( + self, query: str, documents: List[str], top_n: Optional[int] = None + ) -> List[str]: if not top_n: top_n = len(documents) - assert top_n <= len(documents), "'top_n' should be less than or equal to the number of documents" + assert top_n <= len( + documents + ), "'top_n' should be less than or equal to the number of documents" if top_n == 0: return [] url = self.base_url from pyhugegraph.utils.constants import Constants + headers = { "accept": Constants.HEADER_CONTENT_TYPE, "content-type": Constants.HEADER_CONTENT_TYPE, diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py index aa9f0c061..6136d61b4 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/init_reranker.py @@ -32,5 +32,7 @@ def get_reranker(self): model=llm_settings.reranker_model, ) if self.reranker_type == "siliconflow": - return SiliconReranker(api_key=llm_settings.reranker_api_key, model=llm_settings.reranker_model) + return SiliconReranker( + api_key=llm_settings.reranker_api_key, model=llm_settings.reranker_model + ) raise Exception("Reranker type is not supported!") diff --git a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py index d63b0ba3d..e4a9b550a 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py +++ b/hugegraph-llm/src/hugegraph_llm/models/rerankers/siliconflow.py @@ -29,10 +29,14 @@ def __init__( self.api_key = api_key self.model = model - def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int] = None) -> List[str]: + def get_rerank_lists( + self, query: str, documents: List[str], top_n: Optional[int] = None + ) -> List[str]: if not top_n: top_n = len(documents) - assert top_n <= len(documents), "'top_n' should be less than or equal to the number of documents" + assert top_n <= len( + documents + ), "'top_n' should be less than or equal to the number of documents" if top_n == 0: return [] @@ -48,6 +52,7 @@ def get_rerank_lists(self, query: str, documents: List[str], top_n: Optional[int "top_n": top_n, } from pyhugegraph.utils.constants import Constants + headers = { "accept": Constants.HEADER_CONTENT_TYPE, "content-type": Constants.HEADER_CONTENT_TYPE, diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/gremlin_execute.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/gremlin_execute.py new file mode 100644 index 000000000..98fdcdd1b --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/gremlin_execute.py @@ -0,0 +1,68 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from typing import Any, Dict + +from PyCGraph import CStatus + +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.utils.hugegraph_utils import run_gremlin_query + + +def _ensure_limit(query: str, default_limit: int = 100) -> str: + if not query: + return query + q_lower = query.lower() + if "limit(" in q_lower: + return query + if any(token in q_lower for token in ["g.v(", ".v(", "g.e(", ".e("]): + return f"{query}.limit({default_limit})" + return query + + +class GremlinExecuteNode(BaseNode): + def node_init(self): + return CStatus() + + def operator_schedule(self, data_json: Dict[str, Any]): + # Read requested outputs from wk_input + requested = getattr(self.wk_input, "requested_outputs", None) or [] + need_template = "template_execution_result" in requested + need_raw = "raw_execution_result" in requested + + tmpl_q = data_json.get("result", "") + raw_q = data_json.get("raw_result", "") + + if need_template: + try: + safe_q = _ensure_limit(tmpl_q) + data_json["template_exec_res"] = run_gremlin_query(query=safe_q) + except Exception as exc: # pylint: disable=broad-except + data_json["template_exec_res"] = f"{exc}" + else: + data_json["template_exec_res"] = "" + + if need_raw: + try: + safe_q = _ensure_limit(raw_q) + data_json["raw_exec_res"] = run_gremlin_query(query=safe_q) + except Exception as exc: # pylint: disable=broad-except + data_json["raw_exec_res"] = f"{exc}" + else: + data_json["raw_exec_res"] = "" + + return data_json 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 71c490b20..84719d9eb 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py @@ -62,7 +62,7 @@ def node_init(self): return CStatus() def operator_schedule(self, data_json): - print(f"check data json {data_json}") + log.debug("SchemaNode input state: %s", data_json) if self.schema.startswith("{"): try: return self.check_schema.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py new file mode 100644 index 000000000..eb033d869 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -0,0 +1,49 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from typing import Any, Dict + +from PyCGraph import CStatus + +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.gremlin_example_index_query import GremlinExampleIndexQuery +from hugegraph_llm.models.embeddings.init_embedding import Embeddings + + +class GremlinExampleIndexQueryNode(BaseNode): + operator: GremlinExampleIndexQuery + + def node_init(self): + # Build operator (index lazy-loading handled in operator) + embedding = Embeddings().get_embedding() + example_num = getattr(self.wk_input, "example_num", None) + if not isinstance(example_num, int): + example_num = 2 + # Clamp to [0, 10] + example_num = max(0, min(10, example_num)) + self.operator = GremlinExampleIndexQuery(embedding=embedding, num_examples=example_num) + return CStatus() + + def operator_schedule(self, data_json: Dict[str, Any]): + # Ensure query is present in context; degrade gracefully if empty + query = getattr(self.wk_input, "query", "") or "" + data_json["query"] = query + if not query: + data_json["match_result"] = [] + return data_json + # Operator.run writes match_result into context + return self.operator.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py index a28b41346..7df2e68e7 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py @@ -61,16 +61,12 @@ def node_init(self): # few_shot_schema: already parsed dict or raw JSON string few_shot_schema = {} - fss_src = ( - self.wk_input.few_shot_schema if self.wk_input.few_shot_schema else None - ) + fss_src = self.wk_input.few_shot_schema if self.wk_input.few_shot_schema else None if fss_src: try: few_shot_schema = json.loads(fss_src) except json.JSONDecodeError as e: - return CStatus( - -1, f"Few Shot Schema is not in a valid JSON format: {e}" - ) + return CStatus(-1, f"Few Shot Schema is not in a valid JSON format: {e}") _context_payload = { "raw_texts": raw_texts, diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py new file mode 100644 index 000000000..ffbafbaf4 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py @@ -0,0 +1,70 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import json +from typing import Any, Dict, Optional + +from PyCGraph import CStatus + +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize +from hugegraph_llm.models.llms.init_llm import LLMs +from hugegraph_llm.config import prompt as prompt_cfg + + +def _stable_schema_string(state_json: Dict[str, Any]) -> str: + if "simple_schema" in state_json and state_json["simple_schema"] is not None: + return json.dumps(state_json["simple_schema"], ensure_ascii=False, sort_keys=True) + if "schema" in state_json and state_json["schema"] is not None: + return json.dumps(state_json["schema"], ensure_ascii=False, sort_keys=True) + return "" + + +class Text2GremlinNode(BaseNode): + operator: GremlinGenerateSynthesize + + def node_init(self): + # Select LLM + llm = LLMs().get_text2gql_llm() + # Serialize schema deterministically + state_json = self.context.to_json() + schema_str = _stable_schema_string(state_json) + # Prompt fallback + gremlin_prompt: Optional[str] = getattr(self.wk_input, "gremlin_prompt", None) + if gremlin_prompt is None or not str(gremlin_prompt).strip(): + gremlin_prompt = prompt_cfg.gremlin_generate_prompt + # Keep vertices/properties empty for now + self.operator = GremlinGenerateSynthesize( + llm=llm, + schema=schema_str, + vertices=None, + gremlin_prompt=gremlin_prompt, + ) + return CStatus() + + def operator_schedule(self, data_json: Dict[str, Any]): + # Ensure query exists in context; return empty if not provided + query = getattr(self.wk_input, "query", "") or "" + data_json["query"] = query + if not query: + data_json["result"] = "" + data_json["raw_result"] = "" + return data_json + # increase call count for observability + prev = data_json.get("call_count", 0) or 0 + data_json["call_count"] = prev + 1 + return self.operator.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py index c1c742032..fc729c11e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/check_schema.py @@ -59,12 +59,8 @@ def _validate_schema(self, schema: Dict[str, Any]) -> None: check_type(schema, dict, "Input data is not a dictionary.") if "vertexlabels" not in schema or "edgelabels" not in schema: log_and_raise("Input data does not contain 'vertexlabels' or 'edgelabels'.") - check_type( - schema["vertexlabels"], list, "'vertexlabels' in input data is not a list." - ) - check_type( - schema["edgelabels"], list, "'edgelabels' in input data is not a list." - ) + check_type(schema["vertexlabels"], list, "'vertexlabels' in input data is not a list.") + check_type(schema["edgelabels"], list, "'edgelabels' in input data is not a list.") def _process_property_labels(self, schema: Dict[str, Any]) -> (list, set): property_labels = schema.get("propertykeys", []) @@ -82,19 +78,13 @@ def _process_vertex_labels( for vertex_label in schema["vertexlabels"]: self._validate_vertex_label(vertex_label) properties = vertex_label["properties"] - primary_keys = self._process_keys( - vertex_label, "primary_keys", properties[:1] - ) + primary_keys = self._process_keys(vertex_label, "primary_keys", properties[:1]) if len(primary_keys) == 0: log_and_raise(f"'primary_keys' of {vertex_label['name']} is empty.") vertex_label["primary_keys"] = primary_keys - nullable_keys = self._process_keys( - vertex_label, "nullable_keys", properties[1:] - ) + nullable_keys = self._process_keys(vertex_label, "nullable_keys", properties[1:]) vertex_label["nullable_keys"] = nullable_keys - self._add_missing_properties( - properties, property_labels, property_label_set - ) + self._add_missing_properties(properties, property_labels, property_label_set) def _process_edge_labels( self, schema: Dict[str, Any], property_labels: list, property_label_set: set @@ -102,17 +92,13 @@ def _process_edge_labels( for edge_label in schema["edgelabels"]: self._validate_edge_label(edge_label) properties = edge_label.get("properties", []) - self._add_missing_properties( - properties, property_labels, property_label_set - ) + self._add_missing_properties(properties, property_labels, property_label_set) def _validate_vertex_label(self, vertex_label: Dict[str, Any]) -> None: check_type(vertex_label, dict, "VertexLabel in input data is not a dictionary.") if "name" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'name'.") - check_type( - vertex_label["name"], str, "'name' in vertex_label is not of correct type." - ) + check_type(vertex_label["name"], str, "'name' in vertex_label is not of correct type.") if "properties" not in vertex_label: log_and_raise("VertexLabel in input data does not contain 'properties'.") check_type( @@ -133,9 +119,7 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: log_and_raise( "EdgeLabel in input data does not contain 'name', 'source_label', 'target_label'." ) - check_type( - edge_label["name"], str, "'name' in edge_label is not of correct type." - ) + check_type(edge_label["name"], str, "'name' in edge_label is not of correct type.") check_type( edge_label["source_label"], str, @@ -147,13 +131,9 @@ def _validate_edge_label(self, edge_label: Dict[str, Any]) -> None: "'target_label' in edge_label is not of correct type.", ) - def _process_keys( - self, label: Dict[str, Any], key_type: str, default_keys: list - ) -> list: + def _process_keys(self, label: Dict[str, Any], key_type: str, default_keys: list) -> list: keys = label.get(key_type, default_keys) - check_type( - keys, list, f"'{key_type}' in {label['name']} is not of correct type." - ) + check_type(keys, list, f"'{key_type}' in {label['name']} is not of correct type.") new_keys = [key for key in keys if key in label["properties"]] return new_keys diff --git a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py index 910de20d5..dc5b15e00 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py @@ -126,15 +126,20 @@ def _rerank_with_vertex_degree( reranker = Rerankers().get_reranker() try: vertex_rerank_res = [ - reranker.get_rerank_lists(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list + reranker.get_rerank_lists(query, vertex_degree) + [""] + for vertex_degree in vertex_degree_list ] except requests.exceptions.RequestException as e: - log.warning("Online reranker fails, automatically switches to local bleu method: %s", e) + log.warning( + "Online reranker fails, automatically switches to local bleu method: %s", e + ) self.method = "bleu" self.switch_to_bleu = True if self.method == "bleu": - vertex_rerank_res = [_bleu_rerank(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list] + vertex_rerank_res = [ + _bleu_rerank(query, vertex_degree) + [""] for vertex_degree in vertex_degree_list + ] depth = len(vertex_degree_list) for result in results: @@ -144,7 +149,9 @@ def _rerank_with_vertex_degree( knowledge_with_degree[result] += [""] * (depth - len(knowledge_with_degree[result])) def sort_key(res: str) -> Tuple[int, ...]: - return tuple(vertex_rerank_res[i].index(knowledge_with_degree[res][i]) for i in range(depth)) + return tuple( + vertex_rerank_res[i].index(knowledge_with_degree[res][i]) for i in range(depth) + ) sorted_results = sorted(results, key=sort_key) return sorted_results[:topn] diff --git a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py index a873e19ad..6771a9aab 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/document_op/word_extract.py @@ -56,7 +56,8 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context["keywords"] = keywords from hugegraph_llm.utils.log import log - log.info("KEYWORDS: %s", context['keywords']) + + log.info("KEYWORDS: %s", context["keywords"]) return context def _filter_keywords( diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py index be0ac0ca6..58848f827 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py @@ -184,7 +184,7 @@ def merge_dedup_rerank( method=rerank_method, near_neighbor_first=near_neighbor_first, custom_related_information=custom_related_information, - topk_return_results=topk_return_results + topk_return_results=topk_return_results, ) ) return self @@ -238,7 +238,7 @@ def run(self, **kwargs) -> Dict[str, Any]: """ if len(self._operators) == 0: self.extract_keywords().query_graphdb( - max_graph_items=kwargs.get('max_graph_items') + max_graph_items=kwargs.get("max_graph_items") ).synthesize_answer() context = kwargs 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 9eec04f7f..52626b72b 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 @@ -50,9 +50,7 @@ def run(self, data: dict) -> Dict[str, Any]: if not schema: # TODO: ensure the function works correctly (update the logic later) self.schema_free_mode(data.get("triples", [])) - log.warning( - "Using schema_free mode, could try schema_define mode for better effect!" - ) + log.warning("Using schema_free mode, could try schema_define mode for better effect!") else: self.init_schema_if_need(schema) self.load_into_graph(vertices, edges, schema) @@ -68,9 +66,7 @@ def _set_default_property(self, key, input_properties, property_label_map): # list or set default_value = [] input_properties[key] = default_value - log.warning( - "Property '%s' missing in vertex, set to '%s' for now", key, default_value - ) + log.warning("Property '%s' missing in vertex, set to '%s' for now", key, default_value) def _handle_graph_creation(self, func, *args, **kwargs): try: @@ -84,13 +80,9 @@ 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) - vertex_label_map = { - v_label["name"]: v_label for v_label in schema["vertexlabels"] - } + 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"] - } + property_label_map = {p_label["name"]: p_label for p_label in schema["propertykeys"]} for vertex in vertices: input_label = vertex["label"] @@ -106,9 +98,7 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- vertex_label = vertex_label_map[input_label] primary_keys = vertex_label["primary_keys"] nullable_keys = vertex_label.get("nullable_keys", []) - non_null_keys = [ - key for key in vertex_label["properties"] if key not in nullable_keys - ] + non_null_keys = [key for key in vertex_label["properties"] if key not in nullable_keys] has_problem = False # 2. Handle primary-keys mode vertex @@ -140,9 +130,7 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- # 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 - ) + 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(): @@ -179,9 +167,7 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- 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 - ) + self._handle_graph_creation(self.client.graph().addEdge, label, start, end, properties) def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -205,20 +191,18 @@ def init_schema_if_need(self, schema: dict): source_vertex_label = edge["source_label"] target_vertex_label = edge["target_label"] properties = edge["properties"] - self.schema.edgeLabel(edge_label).sourceLabel( - source_vertex_label - ).targetLabel(target_vertex_label).properties(*properties).nullableKeys( - *properties - ).ifNotExist().create() + self.schema.edgeLabel(edge_label).sourceLabel(source_vertex_label).targetLabel( + target_vertex_label + ).properties(*properties).nullableKeys(*properties).ifNotExist().create() def schema_free_mode(self, 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("name").ifNotExist().create() + self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel("vertex").properties( + "name" + ).ifNotExist().create() self.schema.indexLabel("vertexByName").onV("vertex").by( "name" @@ -278,9 +262,7 @@ def _set_property_data_type(self, property_key, data_type): log.warning("UUID type is not supported, use text instead") property_key.asText() else: - log.error( - "Unknown data type %s for property_key %s", data_type, property_key - ) + log.error("Unknown data type %s for property_key %s", data_type, property_key) def _set_property_cardinality(self, property_key, cardinality): if cardinality == PropertyCardinality.SINGLE: @@ -290,13 +272,9 @@ def _set_property_cardinality(self, property_key, cardinality): elif cardinality == PropertyCardinality.SET: property_key.valueSet() else: - log.error( - "Unknown cardinality %s for property_key %s", cardinality, property_key - ) + 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: + def _check_property_data_type(self, data_type: str, cardinality: str, value) -> bool: if cardinality in ( PropertyCardinality.LIST.value, PropertyCardinality.SET.value, @@ -326,9 +304,7 @@ def _check_single_data_type(self, data_type: str, value) -> bool: if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below - if ( - data_type == PropertyDataType.DATE.value - ): # the format should be "yyyy-MM-dd" + if data_type == PropertyDataType.DATE.value: # the format should be "yyyy-MM-dd" import re return isinstance(value, str) and re.match(r"^\d{4}-\d{2}-\d{2}$", value) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py index 6012b7534..bcff5f07b 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py @@ -196,8 +196,8 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: log.debug("Kneighbor gremlin query: %s", gremlin_query) paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = self._format_graph_query_result( - query_paths=paths + graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( + self._format_graph_query_result(query_paths=paths) ) # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) @@ -220,17 +220,21 @@ def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: max_deep=self._max_deep, max_items=self._max_items, ) - log.warning("Unable to find vid, downgraded to property query, please confirm if it meets expectation.") + log.warning( + "Unable to find vid, downgraded to property query, please confirm if it meets expectation." + ) paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = self._format_graph_query_result( - query_paths=paths + graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( + self._format_graph_query_result(query_paths=paths) ) context["graph_result"] = list(graph_chain_knowledge) if context["graph_result"]: context["graph_result_flag"] = 0 - context["vertex_degree_list"] = [list(vertex_degree) for vertex_degree in vertex_degree_list] + context["vertex_degree_list"] = [ + list(vertex_degree) for vertex_degree in vertex_degree_list + ] context["knowledge_with_degree"] = knowledge_with_degree context["graph_context_head"] = ( f"The following are graph knowledge in {self._max_deep} depth, e.g:\n" @@ -272,7 +276,9 @@ def _format_graph_from_vertex(self, query_result: List[Any]) -> Set[str]: knowledge.add(node_str) return knowledge - def _format_graph_query_result(self, query_paths) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: + def _format_graph_query_result( + self, query_paths + ) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: use_id_to_match = self._prop_to_match is None subgraph = set() subgraph_with_degree = {} @@ -282,7 +288,9 @@ def _format_graph_query_result(self, query_paths) -> Tuple[Set[str], List[Set[st for path in query_paths: # 1. Process each path - path_str, vertex_with_degree = self._process_path(path, use_id_to_match, v_cache, e_cache) + path_str, vertex_with_degree = self._process_path( + path, use_id_to_match, v_cache, e_cache + ) subgraph.add(path_str) subgraph_with_degree[path_str] = vertex_with_degree # 2. Update vertex degree list @@ -291,7 +299,11 @@ def _format_graph_query_result(self, query_paths) -> Tuple[Set[str], List[Set[st return subgraph, vertex_degree_list, subgraph_with_degree def _process_path( - self, path: Any, use_id_to_match: bool, v_cache: Set[str], e_cache: Set[Tuple[str, str, str]] + self, + path: Any, + use_id_to_match: bool, + v_cache: Set[str], + e_cache: Set[Tuple[str, str, str]], ) -> Tuple[str, List[str]]: flat_rel = "" raw_flat_rel = path["objects"] @@ -306,7 +318,14 @@ def _process_path( if i % 2 == 0: # Process each vertex flat_rel, prior_edge_str_len, depth = self._process_vertex( - item, flat_rel, node_cache, prior_edge_str_len, depth, nodes_with_degree, use_id_to_match, v_cache + item, + flat_rel, + node_cache, + prior_edge_str_len, + depth, + nodes_with_degree, + use_id_to_match, + v_cache, ) else: # Process each edge @@ -333,7 +352,9 @@ def _process_vertex( return flat_rel, prior_edge_str_len, depth node_cache.add(matched_str) - props_str = ", ".join(f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v) + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v + ) # TODO: we may remove label id or replace with label name if matched_str in v_cache: @@ -356,10 +377,14 @@ def _process_edge( use_id_to_match: bool, e_cache: Set[Tuple[str, str, str]], ) -> Tuple[str, int]: - props_str = ", ".join(f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v) + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v + ) props_str = f"{{{props_str}}}" if props_str else "" prev_matched_str = ( - raw_flat_rel[i - 1]["id"] if use_id_to_match else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] + raw_flat_rel[i - 1]["id"] + if use_id_to_match + else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] ) edge_key = (item["inV"], item["label"], item["outV"]) @@ -369,12 +394,16 @@ def _process_edge( else: edge_label = item["label"] - edge_str = f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" + edge_str = ( + f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" + ) path_str += edge_str prior_edge_str_len = len(edge_str) return path_str, prior_edge_str_len - def _update_vertex_degree_list(self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str]) -> None: + def _update_vertex_degree_list( + self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str] + ) -> None: for depth, node_str in enumerate(nodes_with_degree): if depth >= len(vertex_degree_list): vertex_degree_list.append(set()) @@ -384,8 +413,8 @@ def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: schema = self._get_graph_schema() vertex_props_str, edge_props_str = schema.split("\n")[:2] # TODO: rename to vertex (also need update in the schema) - vertex_props_str = vertex_props_str[len("Vertex properties: "):].strip("[").strip("]") - edge_props_str = edge_props_str[len("Edge properties: "):].strip("[").strip("]") + vertex_props_str = vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") + edge_props_str = edge_props_str[len("Edge properties: ") :].strip("[").strip("]") vertex_labels = self._extract_label_names(vertex_props_str) edge_labels = self._extract_label_names(edge_props_str) return vertex_labels, edge_labels 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 c4e2124c3..90f1c00ea 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 @@ -40,9 +40,7 @@ def simple_schema(self, schema: Dict[str, Any]) -> Dict[str, Any]: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: new_vertex = { - key: vertex[key] - for key in ["id", "name", "properties"] - if key in vertex + key: vertex[key] for key in ["id", "name", "properties"] if key in vertex } mini_schema["vertexlabels"].append(new_vertex) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py index 657baf68e..6d9f96214 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py @@ -23,17 +23,25 @@ from hugegraph_llm.config import resource_path, llm_settings, huge_settings from hugegraph_llm.indices.vector_index import VectorIndex from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel, get_filename_prefix, get_index_folder_name +from hugegraph_llm.utils.embedding_utils import ( + get_embeddings_parallel, + get_filename_prefix, + get_index_folder_name, +) # FIXME: we need keep the logic same with build_semantic_index.py class BuildGremlinExampleIndex: def __init__(self, embedding: BaseEmbedding, examples: List[Dict[str, str]]): - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "gremlin_examples")) self.examples = examples self.embedding = embedding - self.filename_prefix = get_filename_prefix(llm_settings.embedding_type, getattr(embedding, "model_name", None)) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(embedding, "model_name", None) + ) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: # !: We have assumed that self.example is not empty 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 4b7c4e3d4..5689a59ac 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 @@ -24,15 +24,23 @@ from hugegraph_llm.indices.vector_index import VectorIndex from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager -from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel, get_filename_prefix, get_index_folder_name +from hugegraph_llm.utils.embedding_utils import ( + get_embeddings_parallel, + get_filename_prefix, + get_index_folder_name, +) from hugegraph_llm.utils.log import log class BuildSemanticIndex: def __init__(self, embedding: BaseEmbedding): - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) - self.filename_prefix = get_filename_prefix(llm_settings.embedding_type, getattr(embedding, "model_name", None)) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(embedding, "model_name", None) + ) self.vid_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) self.embedding = embedding self.sm = SchemaManager(huge_settings.graph_name) @@ -42,27 +50,19 @@ def _extract_names(self, vertices: list[str]) -> list[str]: def run(self, context: Dict[str, Any]) -> Dict[str, Any]: vertexlabels = self.sm.schema.getSchema()["vertexlabels"] - all_pk_flag = all( - data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels - ) + all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) past_vids = self.vid_index.properties # TODO: We should build vid vector index separately, especially when the vertices may be very large - present_vids = context[ - "vertices" - ] # Warning: data truncated by fetch_graph_data.py + present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py removed_vids = set(past_vids) - set(present_vids) removed_num = self.vid_index.remove(removed_vids) added_vids = list(set(present_vids) - set(past_vids)) if added_vids: - vids_to_process = ( - self._extract_names(added_vids) if all_pk_flag else added_vids - ) - added_embeddings = asyncio.run( - get_embeddings_parallel(self.embedding, vids_to_process) - ) + vids_to_process = self._extract_names(added_vids) if all_pk_flag else added_vids + added_embeddings = asyncio.run(get_embeddings_parallel(self.embedding, 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.to_index_file(self.index_dir, self.filename_prefix) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py index 5cdad0316..f5fb823c5 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/build_vector_index.py @@ -41,9 +41,7 @@ def __init__(self, embedding: BaseEmbedding): self.filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(self.embedding, "model_name", None) ) - self.vector_index = VectorIndex.from_index_file( - self.index_dir, self.filename_prefix - ) + self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: if "chunks" not in context: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py index 96d1a3833..b680f2ca3 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/gremlin_example_index_query.py @@ -26,7 +26,11 @@ from hugegraph_llm.indices.vector_index import VectorIndex, INDEX_FILE_NAME, PROPERTIES_FILE_NAME from hugegraph_llm.models.embeddings.base import BaseEmbedding from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.utils.embedding_utils import get_embeddings_parallel, get_filename_prefix, get_index_folder_name +from hugegraph_llm.utils.embedding_utils import ( + get_embeddings_parallel, + get_filename_prefix, + get_index_folder_name, +) from hugegraph_llm.utils.log import log @@ -34,16 +38,25 @@ class GremlinExampleIndexQuery: def __init__(self, embedding: BaseEmbedding = None, num_examples: int = 1): self.embedding = embedding or Embeddings().get_embedding() self.num_examples = num_examples - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "gremlin_examples")) - self.filename_prefix = get_filename_prefix(llm_settings.embedding_type, - getattr(self.embedding, "model_name", None)) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(self.embedding, "model_name", None) + ) self._ensure_index_exists() self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) def _ensure_index_exists(self): - index_name = f"{self.filename_prefix}_{INDEX_FILE_NAME}" if self.filename_prefix else INDEX_FILE_NAME - props_name = f"{self.filename_prefix}_{PROPERTIES_FILE_NAME}" if self.filename_prefix else PROPERTIES_FILE_NAME + index_name = ( + f"{self.filename_prefix}_{INDEX_FILE_NAME}" if self.filename_prefix else INDEX_FILE_NAME + ) + props_name = ( + f"{self.filename_prefix}_{PROPERTIES_FILE_NAME}" + if self.filename_prefix + else PROPERTIES_FILE_NAME + ) if not ( os.path.exists(os.path.join(self.index_dir, index_name)) and os.path.exists(os.path.join(self.index_dir, props_name)) @@ -61,7 +74,9 @@ def _get_match_result(self, context: Dict[str, Any], query: str) -> List[Dict[st return self.vector_index.search(query_embedding, self.num_examples, dis_threshold=1.8) def _build_default_example_index(self): - properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict(orient="records") + properties = pd.read_csv(os.path.join(resource_path, "demo", "text2gremlin.csv")).to_dict( + orient="records" + ) # TODO: reuse the logic in build_semantic_index.py (consider extract the batch-embedding method) queries = [row["query"] for row in properties] embeddings = asyncio.run(get_embeddings_parallel(self.embedding, queries)) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py index 8e195453d..3ac03246f 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py @@ -31,16 +31,20 @@ class SemanticIdQuery: ID_QUERY_TEMPL = "g.V({vids_str}).limit(8)" def __init__( - self, - embedding: BaseEmbedding, - by: Literal["query", "keywords"] = "keywords", - topk_per_query: int = 10, - topk_per_keyword: int = huge_settings.topk_per_keyword, - vector_dis_threshold: float = huge_settings.vector_dis_threshold, + self, + embedding: BaseEmbedding, + by: Literal["query", "keywords"] = "keywords", + topk_per_query: int = 10, + topk_per_keyword: int = huge_settings.topk_per_keyword, + vector_dis_threshold: float = huge_settings.vector_dis_threshold, ): - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) - self.filename_prefix = get_filename_prefix(llm_settings.embedding_type, getattr(embedding, "model_name", None)) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(embedding, "model_name", None) + ) self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) self.embedding = embedding self.by = by @@ -65,7 +69,7 @@ def _exact_match_vids(self, keywords: List[str]) -> Tuple[List[str], List[str]]: vids_str = ",".join([f"'{vid}'" for vid in possible_vids]) resp = self._client.gremlin().exec(SemanticIdQuery.ID_QUERY_TEMPL.format(vids_str=vids_str)) - searched_vids = [v['id'] for v in resp['data']] + searched_vids = [v["id"] for v in resp["data"]] unsearched_keywords = set(keywords) for vid in searched_vids: @@ -79,10 +83,13 @@ def _fuzzy_match_vids(self, keywords: List[str]) -> List[str]: fuzzy_match_result = [] for keyword in keywords: keyword_vector = self.embedding.get_texts_embeddings([keyword])[0] - results = self.vector_index.search(keyword_vector, top_k=self.topk_per_keyword, - dis_threshold=float(self.vector_dis_threshold)) + results = self.vector_index.search( + keyword_vector, + top_k=self.topk_per_keyword, + dis_threshold=float(self.vector_dis_threshold), + ) if results: - fuzzy_match_result.extend(results[:self.topk_per_keyword]) + fuzzy_match_result.extend(results[: self.topk_per_keyword]) return fuzzy_match_result def run(self, context: Dict[str, Any]) -> Dict[str, Any]: @@ -92,7 +99,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: query_vector = self.embedding.get_texts_embeddings([query])[0] results = self.vector_index.search(query_vector, top_k=self.topk_per_query) if results: - graph_query_list.update(results[:self.topk_per_query]) + graph_query_list.update(results[: self.topk_per_query]) else: # by keywords keywords = context.get("keywords", []) if not keywords: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py index e29f50a76..4ed616929 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/vector_index_query.py @@ -30,9 +30,13 @@ class VectorIndexQuery: def __init__(self, embedding: BaseEmbedding, topk: int = 3): self.embedding = embedding self.topk = topk - self.folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + self.folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) self.index_dir = str(os.path.join(resource_path, self.folder_name, "chunks")) - self.filename_prefix = get_filename_prefix(llm_settings.embedding_type, getattr(embedding, "model_name", None)) + self.filename_prefix = get_filename_prefix( + llm_settings.embedding_type, getattr(embedding, "model_name", None) + ) self.vector_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py index 4348477f6..3b5c63103 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py @@ -37,7 +37,12 @@ class KgBuilder: - def __init__(self, llm: BaseLLM, embedding: Optional[BaseEmbedding] = None, graph: Optional[PyHugeClient] = None): + def __init__( + self, + llm: BaseLLM, + embedding: Optional[BaseEmbedding] = None, + graph: Optional[PyHugeClient] = None, + ): self.operators = [] self.llm = llm self.embedding = embedding @@ -69,7 +74,9 @@ def chunk_split( return self def extract_info( - self, example_prompt: Optional[str] = None, extract_type: Literal["triples", "property_graph"] = "triples" + self, + example_prompt: Optional[str] = None, + extract_type: Literal["triples", "property_graph"] = "triples", ): if extract_type == "triples": self.operators.append(InfoExtract(self.llm, example_prompt)) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py index 5c4ab5fd3..9138f9e9b 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/answer_synthesize.py @@ -62,17 +62,26 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = (f"{context_head_str}\n" - f"{self._context_body}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) response = self._llm.generate(prompt=final_prompt) return {"answer": response} graph_result_context, vector_result_context = self.handle_vector_graph(context) - context = asyncio.run(self.async_generate(context, context_head_str, context_tail_str, - vector_result_context, graph_result_context)) + context = asyncio.run( + self.async_generate( + context, + context_head_str, + context_tail_str, + vector_result_context, + graph_result_context, + ) + ) return context def init_llm(self, context): @@ -95,7 +104,9 @@ def handle_vector_graph(self, context): vector_result_context = "No (vector)phrase related to the query." graph_result = context.get("graph_result") if graph_result: - graph_context_head = context.get("graph_context_head", "Knowledge from graphdb for the query:\n") + graph_context_head = context.get( + "graph_context_head", "Knowledge from graphdb for the query:\n" + ) graph_result_context = graph_context_head + "\n".join( f"{i + 1}. {res}" for i, res in enumerate(graph_result) ) @@ -108,11 +119,13 @@ async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[st context_head_str, context_tail_str = self.init_llm(context) if self._context_body is not None: - context_str = (f"{context_head_str}\n" - f"{self._context_body}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" f"{self._context_body}\n" f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) response = self._llm.generate(prompt=final_prompt) yield {"answer": response} return @@ -120,45 +133,60 @@ async def run_streaming(self, context: Dict[str, Any]) -> AsyncGenerator[Dict[st graph_result_context, vector_result_context = self.handle_vector_graph(context) async for context in self.async_streaming_generate( - context, - context_head_str, - context_tail_str, - vector_result_context, - graph_result_context + context, context_head_str, context_tail_str, vector_result_context, graph_result_context ): yield context - async def async_generate(self, context: Dict[str, Any], context_head_str: str, - context_tail_str: str, vector_result_context: str, - graph_result_context: str): + async def async_generate( + self, + context: Dict[str, Any], + context_head_str: str, + context_tail_str: str, + vector_result_context: str, + graph_result_context: str, + ): # async_tasks stores the async tasks for different answer types async_tasks = {} if self._raw_answer: final_prompt = self._question async_tasks["raw_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) if self._vector_only_answer: - context_str = (f"{context_head_str}\n" - f"{vector_result_context}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" + f"{vector_result_context}\n" + f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) - async_tasks["vector_only_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) + async_tasks["vector_only_task"] = asyncio.create_task( + self._llm.agenerate(prompt=final_prompt) + ) if self._graph_only_answer: - context_str = (f"{context_head_str}\n" - f"{graph_result_context}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" + f"{graph_result_context}\n" + f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) - async_tasks["graph_only_task"] = asyncio.create_task(self._llm.agenerate(prompt=final_prompt)) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) + async_tasks["graph_only_task"] = asyncio.create_task( + self._llm.agenerate(prompt=final_prompt) + ) if self._graph_vector_answer: context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = (f"{context_head_str}\n" - f"{context_body_str}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_tasks["graph_vector_task"] = asyncio.create_task( self._llm.agenerate(prompt=final_prompt) ) @@ -167,7 +195,7 @@ async def async_generate(self, context: Dict[str, Any], context_head_str: str, "raw_task": "raw_answer", "vector_only_task": "vector_only_answer", "graph_only_task": "graph_only_answer", - "graph_vector_task": "graph_vector_answer" + "graph_vector_task": "graph_vector_answer", } for task_key, context_key in async_tasks_mapping.items(): @@ -176,66 +204,95 @@ async def async_generate(self, context: Dict[str, Any], context_head_str: str, context[context_key] = response log.debug("Query Answer: %s", response) - ops = sum([self._raw_answer, self._vector_only_answer, self._graph_only_answer, self._graph_vector_answer]) - context['call_count'] = context.get('call_count', 0) + ops + ops = sum( + [ + self._raw_answer, + self._vector_only_answer, + self._graph_only_answer, + self._graph_vector_answer, + ] + ) + context["call_count"] = context.get("call_count", 0) + ops return context - async def async_streaming_generate(self, context: Dict[str, Any], context_head_str: str, - context_tail_str: str, vector_result_context: str, - graph_result_context: str) -> AsyncGenerator[Dict[str, Any], None]: + async def async_streaming_generate( + self, + context: Dict[str, Any], + context_head_str: str, + context_tail_str: str, + vector_result_context: str, + graph_result_context: str, + ) -> AsyncGenerator[Dict[str, Any], None]: # async_tasks stores the async tasks for different answer types async_generators = [] auto_id = 0 if self._raw_answer: final_prompt = self._question async_generators.append( - self.__llm_generate_with_meta_info(task_id=auto_id, target_key="raw_answer", prompt=final_prompt) + self.__llm_generate_with_meta_info( + task_id=auto_id, target_key="raw_answer", prompt=final_prompt + ) ) auto_id += 1 if self._vector_only_answer: - context_str = (f"{context_head_str}\n" - f"{vector_result_context}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" + f"{vector_result_context}\n" + f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_generators.append( self.__llm_generate_with_meta_info( - task_id=auto_id, - target_key="vector_only_answer", - prompt=final_prompt + task_id=auto_id, target_key="vector_only_answer", prompt=final_prompt ) ) auto_id += 1 if self._graph_only_answer: - context_str = (f"{context_head_str}\n" - f"{graph_result_context}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" + f"{graph_result_context}\n" + f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_generators.append( - self.__llm_generate_with_meta_info(task_id=auto_id, target_key="graph_only_answer", prompt=final_prompt) + self.__llm_generate_with_meta_info( + task_id=auto_id, target_key="graph_only_answer", prompt=final_prompt + ) ) auto_id += 1 if self._graph_vector_answer: context_body_str = f"{vector_result_context}\n{graph_result_context}" if context.get("graph_ratio", 0.5) < 0.5: context_body_str = f"{graph_result_context}\n{vector_result_context}" - context_str = (f"{context_head_str}\n" - f"{context_body_str}\n" - f"{context_tail_str}".strip("\n")) + context_str = ( + f"{context_head_str}\n" f"{context_body_str}\n" f"{context_tail_str}".strip("\n") + ) - final_prompt = self._prompt_template.format(context_str=context_str, query_str=self._question) + final_prompt = self._prompt_template.format( + context_str=context_str, query_str=self._question + ) async_generators.append( self.__llm_generate_with_meta_info( - task_id=auto_id, - target_key="graph_vector_answer", - prompt=final_prompt + task_id=auto_id, target_key="graph_vector_answer", prompt=final_prompt ) ) auto_id += 1 - ops = sum([self._raw_answer, self._vector_only_answer, self._graph_only_answer, self._graph_vector_answer]) - context['call_count'] = context.get('call_count', 0) + ops + ops = sum( + [ + self._raw_answer, + self._vector_only_answer, + self._graph_only_answer, + self._graph_vector_answer, + ] + ) + context["call_count"] = context.get("call_count", 0) + ops async_tasks = [asyncio.create_task(anext(gen)) for gen in async_generators] while True: diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py index 817065aa0..2ac2eafff 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/disambiguate_data.py @@ -53,7 +53,8 @@ def run(self, data: Dict) -> Dict[str, List[Any]]: extract_triples_by_regex(llm_output, data) print( f"LLM {self.__class__.__name__} input:{prompt} \n" - f" output: {llm_output} \n data: {data}") + f" output: {llm_output} \n data: {data}" + ) data["call_count"] = data.get("call_count", 0) + 1 return data diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py index 11f0f6022..650834300 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/gremlin_generate.py @@ -54,7 +54,8 @@ def _format_examples(self, examples: Optional[List[Dict[str, str]]]) -> Optional example_strings = [] for example in examples: example_strings.append( - f"- query: {example['query']}\n" f"- gremlin:\n```gremlin\n{example['gremlin']}\n```" + f"- query: {example['query']}\n" + f"- gremlin:\n```gremlin\n{example['gremlin']}\n```" ) return "\n\n".join(example_strings) @@ -89,11 +90,17 @@ async def async_generate(self, context: Dict[str, Any]): vertices=self._format_vertices(vertices=self.vertices), properties=self._format_properties(properties=None), ) - async_tasks["initialized_answer"] = asyncio.create_task(self.llm.agenerate(prompt=init_prompt)) + async_tasks["initialized_answer"] = asyncio.create_task( + self.llm.agenerate(prompt=init_prompt) + ) raw_response = await async_tasks["raw_answer"] initialized_response = await async_tasks["initialized_answer"] - log.debug("Text2Gremlin with tmpl prompt:\n %s,\n LLM Response: %s", init_prompt, initialized_response) + log.debug( + "Text2Gremlin with tmpl prompt:\n %s,\n LLM Response: %s", + init_prompt, + initialized_response, + ) context["result"] = self._extract_response(response=initialized_response) context["raw_result"] = self._extract_response(response=raw_response) @@ -123,7 +130,11 @@ def sync_generate(self, context: Dict[str, Any]): ) initialized_response = self.llm.generate(prompt=init_prompt) - log.debug("Text2Gremlin with tmpl prompt:\n %s,\n LLM Response: %s", init_prompt, initialized_response) + log.debug( + "Text2Gremlin with tmpl prompt:\n %s,\n LLM Response: %s", + init_prompt, + initialized_response, + ) context["result"] = self._extract_response(response=initialized_response) context["raw_result"] = self._extract_response(response=raw_response) 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 571ffde51..8897e0fea 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 @@ -198,12 +198,8 @@ def valid(self, element_id: str, max_length: int = 256) -> bool: return True def _filter_long_id(self, graph) -> Dict[str, List[Any]]: - graph["vertices"] = [ - vertex for vertex in graph["vertices"] if self.valid(vertex["id"]) - ] + graph["vertices"] = [vertex for vertex in graph["vertices"] if self.valid(vertex["id"])] graph["edges"] = [ - edge - for edge in graph["edges"] - if self.valid(edge["start"]) and self.valid(edge["end"]) + edge for edge in graph["edges"] if self.valid(edge["start"]) and self.valid(edge["end"]) ] return graph diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index 1e9ca652b..32ed9651e 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -151,6 +151,7 @@ def _extract_keywords_from_response( response: str, lowercase: bool = True, start_token: str = "", +<<<<<<< HEAD ) -> Dict[str, float]: results = {} @@ -181,4 +182,27 @@ def _extract_keywords_from_response( except (ValueError, AttributeError) as e: log.warning("Failed to parse item '%s': %s", item, e) continue +======= + ) -> Set[str]: + keywords = [] + # use re.escape(start_token) if start_token contains special chars like */&/^ etc. + matches = re.findall(rf"{start_token}[^\n]+\n?", response) + + for match in matches: + match = match[len(start_token) :].strip() + keywords.extend( + k.lower() if lowercase else k + for k in re.split(r"[,，]+", match) + if len(k.strip()) > 1 + ) + + # if the keyword consists of multiple words, split into sub-words (removing stopwords) + results = set(keywords) + for token in keywords: + sub_tokens = re.findall(r"\w+", token) + if len(sub_tokens) > 1: + results.update( + w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language) + ) +>>>>>>> 78011d3 (Refactor: text2germlin with PCgraph framework (#50)) return results diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py index 82326f000..058d1bce9 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/prompt_generate.py @@ -52,11 +52,11 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: few_shot_example = self._load_few_shot_example(example_name) meta_prompt = prompt_tpl.generate_extract_prompt_template.format( - few_shot_text=few_shot_example.get('text', ''), - few_shot_prompt=few_shot_example.get('prompt', ''), + few_shot_text=few_shot_example.get("text", ""), + few_shot_prompt=few_shot_example.get("prompt", ""), user_text=source_text, user_scenario=scenario, - language=prompt_tpl.llm_settings.language + language=prompt_tpl.llm_settings.language, ) log.debug("Meta-prompt sent to LLM: %s", meta_prompt) generated_prompt = self.llm.generate(prompt=meta_prompt) 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 79fb33b4f..565d79023 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 @@ -67,9 +67,9 @@ def filter_item(schema, items) -> List[Dict[str, Any]]: item_type = item["type"] if item_type == "vertex": label = item["label"] - non_nullable_keys = set( - properties_map[item_type][label]["properties"] - ).difference(set(properties_map[item_type][label]["nullable_keys"])) + non_nullable_keys = set(properties_map[item_type][label]["properties"]).difference( + set(properties_map[item_type][label]["nullable_keys"]) + ) for key in non_nullable_keys: if key not in item["properties"]: item["properties"][key] = "NULL" @@ -82,9 +82,7 @@ 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) -> None: self.llm = llm self.example_prompt = example_prompt self.NECESSARY_ITEM_KEYS = {"label", "type", "properties"} # pylint: disable=invalid-name @@ -142,9 +140,7 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: and "vertices" in property_graph and "edges" in property_graph ): - log.critical( - "Invalid property graph format; expecting 'vertices' and 'edges'." - ) + log.critical("Invalid property graph format; expecting 'vertices' and 'edges'.") return items # Create sets for valid vertex and edge labels based on the schema @@ -154,9 +150,7 @@ def _extract_and_filter_label(self, schema, text) -> List[Dict[str, Any]]: def process_items(item_list, valid_labels, item_type): for item in item_list: if not isinstance(item, dict): - log.warning( - "Invalid property graph item type '%s'.", type(item) - ) + log.warning("Invalid property graph item type '%s'.", type(item)) continue if not self.NECESSARY_ITEM_KEYS.issubset(item.keys()): log.warning("Invalid item keys '%s'.", item.keys()) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py index 53587381a..928948413 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/schema_build.py @@ -34,7 +34,9 @@ def __init__( ): self.llm = llm or LLMs().get_chat_llm() # TODO: use a basic format for it - self.schema_prompt = schema_prompt or """ + self.schema_prompt = ( + schema_prompt + or """ You are a Graph Schema Generator for Apache HugeGraph. Based on the following three parts of content, output a Schema JSON that complies with HugeGraph specifications: @@ -53,6 +55,7 @@ def __init__( - Ensure the schema follows HugeGraph specifications - Do not include comments or extra fields. """ + ) def _format_raw_texts(self, raw_texts: List[str]) -> str: return "\n".join([f"- {text}" for text in raw_texts]) @@ -86,18 +89,15 @@ def build_prompt( self, raw_texts: List[str], query_examples: List[Dict[str, str]], - few_shot_schema: Dict[str, Any] + few_shot_schema: Dict[str, Any], ) -> str: return self.schema_prompt.format( raw_texts=self._format_raw_texts(raw_texts), query_examples=self._format_query_examples(query_examples), - few_shot_schema=self._format_few_shot_schema(few_shot_schema) + few_shot_schema=self._format_few_shot_schema(few_shot_schema), ) - def run( - self, - context: Dict[str, Any] - ) -> Dict[str, Any]: + def run(self, context: Dict[str, Any]) -> Dict[str, Any]: """Generate schema from context containing raw_texts, query_examples and few_shot_schema. Args: diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 6d3418c00..f941098b1 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -33,6 +33,11 @@ class WkFlowInput(GParam): source_text: str = None # Original text scenario: str = None # Scenario description example_name: str = None # Example name + # Fields for Text2Gremlin + query: str = None + example_num: int = None + gremlin_prompt: str = None + requested_outputs: Optional[List[str]] = None def reset(self, _: CStatus) -> None: self.texts = None @@ -49,6 +54,11 @@ def reset(self, _: CStatus) -> None: self.source_text = None self.scenario = None self.example_name = None + # Text2Gremlin related configuration + self.query = None + self.example_num = None + self.gremlin_prompt = None + self.requested_outputs = None class WkFlowState(GParam): @@ -66,6 +76,12 @@ class WkFlowState(GParam): keywords_embeddings = None generated_extract_prompt: Optional[str] = None + # Fields for Text2Gremlin results + match_result: Optional[List[dict]] = None + result: Optional[str] = None + raw_result: Optional[str] = None + template_exec_res: Optional[Any] = None + raw_exec_res: Optional[Any] = None def setup(self): self.schema = None @@ -74,7 +90,7 @@ def setup(self): self.edges = None self.vertices = None self.triples = None - self.call_count = None + self.call_count = 0 self.keywords = None self.vector_result = None @@ -82,6 +98,12 @@ def setup(self): self.keywords_embeddings = None self.generated_extract_prompt = None + # Text2Gremlin results reset + self.match_result = [] + self.result = "" + self.raw_result = "" + self.template_exec_res = "" + self.raw_exec_res = "" return CStatus() @@ -94,11 +116,7 @@ def to_json(self): dict: A dictionary containing non-None instance members and their serialized values. """ # Only export instance attributes (excluding methods and class attributes) whose values are not None - return { - k: v - for k, v in self.__dict__.items() - if not k.startswith("_") and v is not None - } + return {k: v for k, v in self.__dict__.items() if not k.startswith("_") and v is not None} # Implement a method that assigns keys from data_json as WkFlowState member variables def assign_from_json(self, data_json: dict): diff --git a/hugegraph-llm/src/hugegraph_llm/utils/anchor.py b/hugegraph-llm/src/hugegraph_llm/utils/anchor.py index d5f687a94..4542a7fd9 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/anchor.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/anchor.py @@ -15,16 +15,17 @@ from pathlib import Path + def get_project_root() -> Path: """ Returns the Path object of the project root directory. - - The function searches for common project root indicators like pyproject.toml + + The function searches for common project root indicators like pyproject.toml or .git directory by traversing up the directory tree from the current file location. - + Returns: Path: The absolute path to the project root directory - + Raises: RuntimeError: If no project root indicators could be found """ diff --git a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py index b07de6f4b..2914c4b28 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/decorators.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/decorators.py @@ -109,6 +109,7 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any: def with_task_id(func: Callable) -> Callable: def wrapper(*args: Any, **kwargs: Any) -> Any: import uuid + task_id = f"task_{str(uuid.uuid4())[:8]}" log.debug("New task created with id: %s", task_id) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py index 55e50eadd..b2f485cea 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/embedding_utils.py @@ -24,7 +24,9 @@ from hugegraph_llm.models.embeddings.base import BaseEmbedding -async def _get_batch_with_progress(embedding: BaseEmbedding, batch: list[str], pbar: tqdm) -> list[Any]: +async def _get_batch_with_progress( + embedding: BaseEmbedding, batch: list[str], pbar: tqdm +) -> list[Any]: result = await embedding.async_get_texts_embeddings(batch) pbar.update(1) return result @@ -58,10 +60,7 @@ async def get_embeddings_parallel(embedding: BaseEmbedding, vids: list[str]) -> embeddings = [] with tqdm(total=len(vid_batches)) as pbar: # Create tasks for each batch with progress bar updates - tasks = [ - _get_batch_with_progress(embedding, batch, pbar) - for batch in vid_batches - ] + tasks = [_get_batch_with_progress(embedding, batch, pbar) for batch in vid_batches] # Use asyncio.gather() to preserve order batch_results = await asyncio.gather(*tasks) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index ccace69f2..7b870033a 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -45,13 +45,9 @@ def get_graph_index_info(): def get_graph_index_info_old(): - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) graph_summary_info = builder.fetch_graph_data().run() - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(builder.embedding, "model_name", None) @@ -66,16 +62,12 @@ def get_graph_index_info_old(): def clean_all_graph_index(): - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None), ) - VectorIndex.clean( - str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix - ) + VectorIndex.clean(str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix) VectorIndex.clean( str(os.path.join(resource_path, folder_name, "gremlin_examples")), filename_prefix, @@ -107,18 +99,14 @@ def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: def extract_graph_origin(input_file, input_text, schema, example_prompt) -> str: texts = read_documents(input_file, input_text) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) if not schema: return "ERROR: please input with correct schema/format." error_message = parse_schema(schema, builder) if error_message: return error_message - builder.chunk_split(texts, "document", "zh").extract_info( - example_prompt, "property_graph" - ) + builder.chunk_split(texts, "document", "zh").extract_info(example_prompt, "property_graph") try: context = builder.run() @@ -168,9 +156,7 @@ def update_vid_embedding(): def update_vid_embedding_old(): - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) builder.fetch_graph_data().build_vertex_id_semantic_index() log.debug("Operators: %s", builder.operators) try: @@ -199,9 +185,7 @@ def import_graph_data_old(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: data_json = json.loads(data.strip()) log.debug("Import graph data: %s", data) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) if schema: error_message = parse_schema(schema, builder) if error_message: @@ -222,9 +206,7 @@ def import_graph_data_old(data: str, schema: str) -> Union[str, Dict[str, Any]]: def build_schema(input_text, query_example, few_shot): scheduler = SchedulerSingleton.get_instance() try: - return scheduler.schedule_flow( - "build_schema", input_text, query_example, few_shot - ) + return scheduler.schedule_flow("build_schema", input_text, query_example, few_shot) except (TypeError, ValueError) as e: raise gr.Error(f"Schema generation failed: {e}") @@ -257,9 +239,7 @@ def build_schema_old(input_text, query_example, few_shot): except json.JSONDecodeError as e: raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) + builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) try: schema = builder.build_schema().run(context) except Exception as e: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py index 1d02b45d3..147c0074c 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py @@ -53,7 +53,9 @@ def init_hg_test_data(): schema = client.schema() schema.propertyKey("name").asText().ifNotExist().create() schema.propertyKey("birthDate").asText().ifNotExist().create() - schema.vertexLabel("Person").properties("name", "birthDate").useCustomizeStringId().ifNotExist().create() + schema.vertexLabel("Person").properties( + "name", "birthDate" + ).useCustomizeStringId().ifNotExist().create() schema.vertexLabel("Movie").properties("name").useCustomizeStringId().ifNotExist().create() schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() @@ -110,13 +112,13 @@ def backup_data(): files = { "vertices.json": f"g.V().limit({MAX_VERTICES})" - f".aggregate('vertices').count().as('count').select('count','vertices')", + f".aggregate('vertices').count().as('count').select('count','vertices')", "edges.json": f"g.E().limit({MAX_EDGES}).aggregate('edges').count().as('count').select('count','edges')", - "schema.json": client.schema().getSchema(_format="groovy") + "schema.json": client.schema().getSchema(_format="groovy"), } vertexlabels = client.schema().getSchema()["vertexlabels"] - all_pk_flag = all(data.get('id_strategy') == 'PRIMARY_KEY' for data in vertexlabels) + all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) for filename, query in files.items(): write_backup_file(client, backup_subdir, filename, query, all_pk_flag) @@ -137,14 +139,22 @@ def write_backup_file(client, backup_subdir, filename, query, all_pk_flag): json.dump(data, f, ensure_ascii=False) elif filename == "vertices.json": data_full = client.gremlin().exec(query)["data"][0]["vertices"] - data = [{key: value for key, value in vertex.items() if key != "id"} - for vertex in data_full] if all_pk_flag else data_full + data = ( + [ + {key: value for key, value in vertex.items() if key != "id"} + for vertex in data_full + ] + if all_pk_flag + else data_full + ) json.dump(data, f, ensure_ascii=False) elif filename == "schema.json": data_full = query if isinstance(data_full, dict) and "schema" in data_full: groovy_filename = filename.replace(".json", ".groovy") - with open(os.path.join(backup_subdir, groovy_filename), "w", encoding="utf-8") as groovy_file: + with open( + os.path.join(backup_subdir, groovy_filename), "w", encoding="utf-8" + ) as groovy_file: groovy_file.write(str(data_full["schema"])) else: data = data_full @@ -171,7 +181,7 @@ def manage_backup_retention(): raise Exception("Failed to manage backup retention") from e -#TODO: In the path demo/rag_demo/configs_block.py, +# TODO: In the path demo/rag_demo/configs_block.py, # there is a function test_api_connection that is similar to this function, # but it is not straightforward to reuse def check_graph_db_connection(url: str, name: str, user: str, pwd: str, graph_space: str) -> bool: diff --git a/hugegraph-llm/src/hugegraph_llm/utils/log.py b/hugegraph-llm/src/hugegraph_llm/utils/log.py index 7076869fd..b64017454 100755 --- a/hugegraph-llm/src/hugegraph_llm/utils/log.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/log.py @@ -31,7 +31,7 @@ log_level=INFO, logger_name="root", propagate_logs=True, - stdout_logging=True + stdout_logging=True, ) # Initialize custom logger diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 138b0d359..301a6bdab 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -50,9 +50,7 @@ def read_documents(input_file, input_text): texts.append(text) elif full_path.endswith(".pdf"): # TODO: support PDF file - raise gr.Error( - "PDF will be supported later! Try to upload text/docx now" - ) + raise gr.Error("PDF will be supported later! Try to upload text/docx now") else: raise gr.Error("Please input txt or docx file.") else: @@ -62,9 +60,7 @@ def read_documents(input_file, input_text): # pylint: disable=C0301 def get_vector_index_info(): - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) ) @@ -91,15 +87,11 @@ def get_vector_index_info(): def clean_vector_index(): - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) + folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) filename_prefix = get_filename_prefix( llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) ) - VectorIndex.clean( - str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix - ) + VectorIndex.clean(str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix) gr.Info("Clean vector index successfully!") diff --git a/hugegraph-llm/src/tests/config/test_config.py b/hugegraph-llm/src/tests/config/test_config.py index 6c803135f..7f480befa 100644 --- a/hugegraph-llm/src/tests/config/test_config.py +++ b/hugegraph-llm/src/tests/config/test_config.py @@ -23,5 +23,6 @@ class TestConfig(unittest.TestCase): def test_config(self): import nltk from hugegraph_llm.config import resource_path + nltk.data.path.append(resource_path) nltk.data.find("corpora/stopwords") diff --git a/hugegraph-llm/src/tests/models/embeddings/test_openai_embedding.py b/hugegraph-llm/src/tests/models/embeddings/test_openai_embedding.py index b9ded0f6c..f7afd15c6 100644 --- a/hugegraph-llm/src/tests/models/embeddings/test_openai_embedding.py +++ b/hugegraph-llm/src/tests/models/embeddings/test_openai_embedding.py @@ -22,6 +22,7 @@ class TestOpenAIEmbedding(unittest.TestCase): def test_embedding_dimension(self): from hugegraph_llm.models.embeddings.openai import OpenAIEmbedding + embedding = OpenAIEmbedding(api_key="") result = embedding.get_text_embedding("hello world!") print(result) diff --git a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py index caabe2a8e..7ad914468 100644 --- a/hugegraph-llm/src/tests/models/llms/test_ollama_client.py +++ b/hugegraph-llm/src/tests/models/llms/test_ollama_client.py @@ -28,7 +28,10 @@ def test_generate(self): def test_stream_generate(self): ollama_client = OllamaClient(model="llama3:8b-instruct-fp16") + def on_token_callback(chunk): print(chunk, end="", flush=True) - ollama_client.generate_streaming(prompt="What is the capital of France?", - on_token_callback=on_token_callback) + + ollama_client.generate_streaming( + prompt="What is the capital of France?", on_token_callback=on_token_callback + ) diff --git a/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py b/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py index d20a198f2..317d02879 100644 --- a/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py +++ b/hugegraph-llm/src/tests/operators/common_op/test_check_schema.py @@ -26,12 +26,7 @@ def setUp(self): def test_schema_check_with_valid_input(self): data = { - "vertexlabels": [ - { - "name": "person", - "properties": ["name", "age", "occupation"] - } - ], + "vertexlabels": [{"name": "person", "properties": ["name", "age", "occupation"]}], "edgelabels": [ { "name": "knows", @@ -41,7 +36,7 @@ def test_schema_check_with_valid_input(self): ], } check_schema = CheckSchema(data) - self.assertEqual(check_schema.run(), {'schema': data}) + self.assertEqual(check_schema.run(), {"schema": data}) def test_schema_check_with_invalid_input(self): data = "invalid input" diff --git a/hugegraph-llm/src/tests/operators/common_op/test_nltk_helper.py b/hugegraph-llm/src/tests/operators/common_op/test_nltk_helper.py index 5ad73ed6f..b557cfc1b 100644 --- a/hugegraph-llm/src/tests/operators/common_op/test_nltk_helper.py +++ b/hugegraph-llm/src/tests/operators/common_op/test_nltk_helper.py @@ -22,6 +22,7 @@ class TestNLTKHelper(unittest.TestCase): def test_stopwords(self): from hugegraph_llm.operators.common_op.nltk_helper import NLTKHelper + nltk_helper = NLTKHelper() stopwords = nltk_helper.stopwords() print(stopwords) diff --git a/hugegraph-python-client/src/pyhugegraph/api/auth.py b/hugegraph-python-client/src/pyhugegraph/api/auth.py index 90b3e98d0..ab7d66169 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/auth.py +++ b/hugegraph-python-client/src/pyhugegraph/api/auth.py @@ -84,9 +84,7 @@ def create_group(self, group_name, group_description=None) -> Optional[Dict]: return self._invoke_request(data=json.dumps(data)) @router.http("DELETE", "auth/groups/{group_id}") - def delete_group( - self, group_id # pylint: disable=unused-argument - ) -> Optional[Dict]: + def delete_group(self, group_id) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("PUT", "auth/groups/{group_id}") @@ -116,9 +114,7 @@ def grant_accesses(self, group_id, target_id, access_permission) -> Optional[Dic ) @router.http("DELETE", "auth/accesses/{access_id}") - def revoke_accesses( - self, access_id # pylint: disable=unused-argument - ) -> Optional[Dict]: + def revoke_accesses(self, access_id) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("PUT", "auth/accesses/{access_id}") @@ -130,9 +126,7 @@ def modify_accesses( return self._invoke_request(data=json.dumps(data)) @router.http("GET", "auth/accesses/{access_id}") - def get_accesses( - self, access_id # pylint: disable=unused-argument - ) -> Optional[Dict]: + def get_accesses(self, access_id) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "auth/accesses") @@ -205,9 +199,7 @@ def update_belong( return self._invoke_request(data=json.dumps(data)) @router.http("GET", "auth/belongs/{belong_id}") - def get_belong( - self, belong_id # pylint: disable=unused-argument - ) -> Optional[Dict]: + def get_belong(self, belong_id) -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "auth/belongs") diff --git a/hugegraph-python-client/src/pyhugegraph/api/graph.py b/hugegraph-python-client/src/pyhugegraph/api/graph.py index 907e01a5b..4555eeda4 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/graph.py +++ b/hugegraph-python-client/src/pyhugegraph/api/graph.py @@ -141,9 +141,7 @@ def addEdges(self, input_data) -> Optional[List[EdgeData]]: def appendEdge( self, edge_id, properties # pylint: disable=unused-argument ) -> Optional[EdgeData]: - if response := self._invoke_request( - data=json.dumps({"properties": properties}) - ): + if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None @@ -151,16 +149,12 @@ def appendEdge( def eliminateEdge( self, edge_id, properties # pylint: disable=unused-argument ) -> Optional[EdgeData]: - if response := self._invoke_request( - data=json.dumps({"properties": properties}) - ): + if response := self._invoke_request(data=json.dumps({"properties": properties})): return EdgeData(response) return None @router.http("GET", "graph/edges/{edge_id}") - def getEdgeById( - self, edge_id # pylint: disable=unused-argument - ) -> Optional[EdgeData]: + def getEdgeById(self, edge_id) -> Optional[EdgeData]: # pylint: disable=unused-argument if response := self._invoke_request(): return EdgeData(response) return None diff --git a/hugegraph-python-client/src/pyhugegraph/api/schema.py b/hugegraph-python-client/src/pyhugegraph/api/schema.py index 8b4f54cfe..7e8926678 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/schema.py +++ b/hugegraph-python-client/src/pyhugegraph/api/schema.py @@ -64,9 +64,7 @@ def indexLabel(self, name): return index_label @router.http("GET", "schema?format={_format}") - def getSchema( - self, _format: str = "json" # pylint: disable=unused-argument - ) -> Optional[Dict]: + def getSchema(self, _format: str = "json") -> Optional[Dict]: # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "schema/propertykeys/{property_name}") @@ -84,9 +82,7 @@ def getPropertyKeys(self) -> Optional[List[PropertyKeyData]]: return None @router.http("GET", "schema/vertexlabels/{name}") - def getVertexLabel( - self, name # pylint: disable=unused-argument - ) -> Optional[VertexLabelData]: + def getVertexLabel(self, name) -> Optional[VertexLabelData]: # pylint: disable=unused-argument if response := self._invoke_request(): return VertexLabelData(response) log.error("VertexLabel not found: %s", str(response)) @@ -128,9 +124,7 @@ def getRelations(self) -> Optional[List[str]]: return None @router.http("GET", "schema/indexlabels/{name}") - def getIndexLabel( - self, name # pylint: disable=unused-argument - ) -> Optional[IndexLabelData]: + def getIndexLabel(self, name) -> Optional[IndexLabelData]: # pylint: disable=unused-argument if response := self._invoke_request(): return IndexLabelData(response) log.error("IndexLabel not found: %s", str(response)) diff --git a/hugegraph-python-client/src/pyhugegraph/api/schema_manage/index_label.py b/hugegraph-python-client/src/pyhugegraph/api/schema_manage/index_label.py index 252d487bd..acef8f968 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/schema_manage/index_label.py +++ b/hugegraph-python-client/src/pyhugegraph/api/schema_manage/index_label.py @@ -83,11 +83,13 @@ def ifNotExist(self) -> "IndexLabel": @decorator_create def create(self): dic = self._parameter_holder.get_dic() - data = {"name": dic["name"], - "base_type": dic["base_type"], - "base_value": dic["base_value"], - "index_type": dic["index_type"], - "fields": list(dic["fields"])} + data = { + "name": dic["name"], + "base_type": dic["base_type"], + "base_value": dic["base_value"], + "index_type": dic["index_type"], + "fields": list(dic["fields"]), + } path = "schema/indexlabels" self.clean_parameter_holder() if response := self._sess.request(path, "POST", data=json.dumps(data)): diff --git a/hugegraph-python-client/src/pyhugegraph/api/services.py b/hugegraph-python-client/src/pyhugegraph/api/services.py index e086ae13e..f353673db 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/services.py +++ b/hugegraph-python-client/src/pyhugegraph/api/services.py @@ -87,9 +87,7 @@ def list_services(self, graphspace: str): # pylint: disable=unused-argument return self._invoke_request() @router.http("GET", "/graphspaces/{graphspace}/services/{service}") - def get_service( - self, graphspace: str, service: str # pylint: disable=unused-argument - ): + def get_service(self, graphspace: str, service: str): # pylint: disable=unused-argument """ Retrieve the details of a specific service. @@ -112,9 +110,7 @@ def get_service( """ return self._invoke_request() - def delete_service( - self, graphspace: str, service: str # pylint: disable=unused-argument - ): + def delete_service(self, graphspace: str, service: str): # pylint: disable=unused-argument """ Delete a specific service within a graph space. diff --git a/hugegraph-python-client/src/pyhugegraph/api/traverser.py b/hugegraph-python-client/src/pyhugegraph/api/traverser.py index 628c3f4bd..72dddb07a 100644 --- a/hugegraph-python-client/src/pyhugegraph/api/traverser.py +++ b/hugegraph-python-client/src/pyhugegraph/api/traverser.py @@ -26,33 +26,23 @@ class TraverserManager(HugeParamsBase): def k_out(self, source_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() - @router.http( - "GET", 'traversers/kneighbor?source="{source_id}"&max_depth={max_depth}' - ) + @router.http("GET", 'traversers/kneighbor?source="{source_id}"&max_depth={max_depth}') def k_neighbor(self, source_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() - @router.http( - "GET", 'traversers/sameneighbors?vertex="{vertex_id}"&other="{other_id}"' - ) + @router.http("GET", 'traversers/sameneighbors?vertex="{vertex_id}"&other="{other_id}"') def same_neighbors(self, vertex_id, other_id): # pylint: disable=unused-argument return self._invoke_request() - @router.http( - "GET", 'traversers/jaccardsimilarity?vertex="{vertex_id}"&other="{other_id}"' - ) - def jaccard_similarity( - self, vertex_id, other_id # pylint: disable=unused-argument - ): + @router.http("GET", 'traversers/jaccardsimilarity?vertex="{vertex_id}"&other="{other_id}"') + def jaccard_similarity(self, vertex_id, other_id): # pylint: disable=unused-argument return self._invoke_request() @router.http( "GET", 'traversers/shortestpath?source="{source_id}"&target="{target_id}"&max_depth={max_depth}', ) - def shortest_path( - self, source_id, target_id, max_depth # pylint: disable=unused-argument - ): + def shortest_path(self, source_id, target_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() @router.http( @@ -78,9 +68,7 @@ def weighted_shortest_path( "GET", 'traversers/singlesourceshortestpath?source="{source_id}"&max_depth={max_depth}', ) - def single_source_shortest_path( - self, source_id, max_depth # pylint: disable=unused-argument - ): + def single_source_shortest_path(self, source_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() @router.http("POST", "traversers/multinodeshortestpath") @@ -114,9 +102,17 @@ def multi_node_shortest_path( def paths(self, source_id, target_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() - @router.http("POST", 'traversers/paths') + @router.http("POST", "traversers/paths") def advanced_paths( - self, sources, targets, step, max_depth, nearest=True, capacity=10000000, limit=10, with_vertex=False + self, + sources, + targets, + step, + max_depth, + nearest=True, + capacity=10000000, + limit=10, + with_vertex=False, ): return self._invoke_request( data=json.dumps( @@ -133,7 +129,6 @@ def advanced_paths( ) ) - @router.http("POST", "traversers/customizedpaths") def customized_paths( self, sources, steps, sort_by="INCR", with_vertex=True, capacity=-1, limit=-1 @@ -152,9 +147,7 @@ def customized_paths( ) @router.http("POST", "traversers/templatepaths") - def template_paths( - self, sources, targets, steps, capacity=10000, limit=10, with_vertex=True - ): + def template_paths(self, sources, targets, steps, capacity=10000, limit=10, with_vertex=True): return self._invoke_request( data=json.dumps( { @@ -172,9 +165,7 @@ def template_paths( "GET", 'traversers/crosspoints?source="{source_id}"&target="{target_id}"&max_depth={max_depth}', ) - def crosspoints( - self, source_id, target_id, max_depth # pylint: disable=unused-argument - ): + def crosspoints(self, source_id, target_id, max_depth): # pylint: disable=unused-argument return self._invoke_request() @router.http("POST", "traversers/customizedcrosspoints") diff --git a/hugegraph-python-client/src/pyhugegraph/client.py b/hugegraph-python-client/src/pyhugegraph/client.py index 3b0301321..c9f4d1027 100644 --- a/hugegraph-python-client/src/pyhugegraph/client.py +++ b/hugegraph-python-client/src/pyhugegraph/client.py @@ -53,7 +53,7 @@ def __init__( user: str, pwd: str, graphspace: Optional[str] = None, - timeout: Optional[tuple[float, float]] = None + timeout: Optional[tuple[float, float]] = None, ): self.cfg = HGraphConfig(url, user, pwd, graph, graphspace, timeout or (0.5, 15.0)) diff --git a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py index 4bb70dba5..d5cc0eb9d 100644 --- a/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py +++ b/hugegraph-python-client/src/pyhugegraph/example/hugegraph_example.py @@ -26,15 +26,13 @@ schema = client.schema() schema.propertyKey("name").asText().ifNotExist().create() schema.propertyKey("birthDate").asText().ifNotExist().create() - schema.vertexLabel("Person").properties( - "name", "birthDate" - ).usePrimaryKeyId().primaryKeys("name").ifNotExist().create() - schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys( + schema.vertexLabel("Person").properties("name", "birthDate").usePrimaryKeyId().primaryKeys( "name" ).ifNotExist().create() - schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel( - "Movie" + schema.vertexLabel("Movie").properties("name").usePrimaryKeyId().primaryKeys( + "name" ).ifNotExist().create() + schema.edgeLabel("ActedIn").sourceLabel("Person").targetLabel("Movie").ifNotExist().create() print(schema.getVertexLabels()) print(schema.getEdgeLabels()) @@ -47,9 +45,7 @@ p2 = g.addVertex("Person", {"name": "Robert De Niro", "birthDate": "1943-08-17"}) m1 = g.addVertex("Movie", {"name": "The Godfather"}) m2 = g.addVertex("Movie", {"name": "The Godfather Part II"}) - m3 = g.addVertex( - "Movie", {"name": "The Godfather Coda The Death of Michael Corleone"} - ) + m3 = g.addVertex("Movie", {"name": "The Godfather Coda The Death of Michael Corleone"}) # add Edge g.addEdge("ActedIn", p1.id, m1.id, {}) diff --git a/hugegraph-python-client/src/pyhugegraph/structure/property_key_data.py b/hugegraph-python-client/src/pyhugegraph/structure/property_key_data.py index 6fb7c36f0..ff50d9b2f 100644 --- a/hugegraph-python-client/src/pyhugegraph/structure/property_key_data.py +++ b/hugegraph-python-client/src/pyhugegraph/structure/property_key_data.py @@ -62,5 +62,7 @@ def userdata(self): return self.__user_data def __repr__(self): - res = f"name: {self.__name}, cardinality: {self.__cardinality}, data_type: {self.__data_type}" + res = ( + f"name: {self.__name}, cardinality: {self.__cardinality}, data_type: {self.__data_type}" + ) return res diff --git a/hugegraph-python-client/src/pyhugegraph/utils/huge_config.py b/hugegraph-python-client/src/pyhugegraph/utils/huge_config.py index 3f6d78b95..429c07c6b 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/huge_config.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/huge_config.py @@ -39,7 +39,7 @@ class HGraphConfig: def __post_init__(self): # Add URL prefix compatibility check - if self.url and not self.url.startswith('http'): + if self.url and not self.url.startswith("http"): self.url = f"http://{self.url}" if self.graphspace and self.graphspace.strip(): @@ -47,9 +47,7 @@ def __post_init__(self): else: try: - response = requests.get( - f"{self.url}/versions", timeout=0.5 - ) + response = requests.get(f"{self.url}/versions", timeout=0.5) core = response.json()["versions"]["core"] log.info( # pylint: disable=logging-fstring-interpolation f"Retrieved API version information from the server: {core}." @@ -71,4 +69,6 @@ def __post_init__(self): except Exception: # pylint: disable=broad-exception-caught exc_type, exc_value, tb = sys.exc_info() traceback.print_exception(exc_type, exc_value, tb) - log.warning("Failed to retrieve API version information from the server, reverting to default v1.") + log.warning( + "Failed to retrieve API version information from the server, reverting to default v1." + ) diff --git a/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py b/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py index 0db81ed1c..f4a38a418 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/huge_router.py @@ -81,9 +81,7 @@ def wrapper(self: "HGraphContext", *args: Any, **kwargs: Any) -> Any: route = RouterRegistry().routers.get(func.__qualname__) if route.request_func is None: - route.request_func = functools.partial( - self.session.request, method=method - ) + route.request_func = functools.partial(self.session.request, method=method) return func(self, *args, **kwargs) @@ -134,9 +132,7 @@ def wrapper(self: "HGraphContext", *args: Any, **kwargs: Any) -> Any: formatted_path = path # Use functools.partial to create a partial function for making requests - make_request = functools.partial( - self.session.request, formatted_path, method - ) + make_request = functools.partial(self.session.request, formatted_path, method) # Store the partial function on the instance setattr(self, f"_{func.__name__}_request", make_request) diff --git a/hugegraph-python-client/src/pyhugegraph/utils/log.py b/hugegraph-python-client/src/pyhugegraph/utils/log.py index b263d32d5..c6f6bd074 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/log.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/log.py @@ -138,7 +138,9 @@ def init_logger( def _cached_log_file(filename): """Cache the opened file object""" # Use 1K buffer if writing to cloud storage - with open(filename, "a", buffering=_determine_buffer_size(filename), encoding="utf-8") as file_io: + with open( + filename, "a", buffering=_determine_buffer_size(filename), encoding="utf-8" + ) as file_io: atexit.register(file_io.close) return file_io diff --git a/hugegraph-python-client/src/pyhugegraph/utils/util.py b/hugegraph-python-client/src/pyhugegraph/utils/util.py index 90f27c24a..56a135547 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/util.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/util.py @@ -44,7 +44,9 @@ def create_exception(response_content): def check_if_authorized(response): if response.status_code == 401: - raise NotAuthorizedError(f"Please check your username and password. {str(response.content)}") + raise NotAuthorizedError( + f"Please check your username and password. {str(response.content)}" + ) return True @@ -56,8 +58,12 @@ def check_if_success(response, error=None): req = response.request req_body = req.body if req.body else "Empty body" response_body = response.text if response.text else "Empty body" - log.error("Error-Client: Request URL: %s, Request Body: %s, Response Body: %s", - req.url, req_body, response_body) + log.error( + "Error-Client: Request URL: %s, Request Body: %s, Response Body: %s", + req.url, + req_body, + response_body, + ) raise error return True @@ -103,9 +109,14 @@ def __call__(self, response: requests.Response, method: str, path: str): details = "key 'exception' not found" req_body = response.request.body if response.request.body else "Empty body" - req_body = req_body.encode('utf-8').decode('unicode_escape') - log.error("%s: %s\n[Body]: %s\n[Server Exception]: %s", - method, str(e).encode('utf-8').decode('unicode_escape'), req_body, details) + req_body = req_body.encode("utf-8").decode("unicode_escape") + log.error( + "%s: %s\n[Body]: %s\n[Server Exception]: %s", + method, + str(e).encode("utf-8").decode("unicode_escape"), + req_body, + details, + ) if response.status_code == 404: raise NotFoundError(response.content) from e diff --git a/hugegraph-python-client/src/tests/api/test_auth.py b/hugegraph-python-client/src/tests/api/test_auth.py index d2d30cf26..10e6bad7f 100644 --- a/hugegraph-python-client/src/tests/api/test_auth.py +++ b/hugegraph-python-client/src/tests/api/test_auth.py @@ -98,9 +98,7 @@ def test_group_operations(self): self.assertEqual(group["group_name"], "test_group") # Modify the group - group = self.auth.modify_group( - group["id"], group_description="test_description" - ) + group = self.auth.modify_group(group["id"], group_description="test_description") self.assertEqual(group["group_description"], "test_description") # Delete the group @@ -135,9 +133,7 @@ def test_target_operations(self): [{"type": "VERTEX", "label": "person", "properties": {"city": "Shanghai"}}], ) # Verify the target was modified - self.assertEqual( - target["target_resources"][0]["properties"]["city"], "Shanghai" - ) + self.assertEqual(target["target_resources"][0]["properties"]["city"], "Shanghai") # Delete the target self.auth.delete_target(target["id"]) diff --git a/hugegraph-python-client/src/tests/api/test_version.py b/hugegraph-python-client/src/tests/api/test_version.py index 1d6325dfd..44c5f376c 100644 --- a/hugegraph-python-client/src/tests/api/test_version.py +++ b/hugegraph-python-client/src/tests/api/test_version.py @@ -42,7 +42,7 @@ def tearDown(self): def test_version(self): version = self.version.version() self.assertIsInstance(version, dict) - self.assertIn("version", version['versions']) - self.assertIn("core", version['versions']) - self.assertIn("gremlin", version['versions']) - self.assertIn("api", version['versions']) + self.assertIn("version", version["versions"]) + self.assertIn("core", version["versions"]) + self.assertIn("gremlin", version["versions"]) + self.assertIn("api", version["versions"]) diff --git a/hugegraph-python-client/src/tests/client_utils.py b/hugegraph-python-client/src/tests/client_utils.py index 63b6d0770..f711072b8 100644 --- a/hugegraph-python-client/src/tests/client_utils.py +++ b/hugegraph-python-client/src/tests/client_utils.py @@ -28,7 +28,11 @@ class ClientUtils: def __init__(self): self.client = PyHugeClient( - url=self.URL, user=self.USERNAME, pwd=self.PASSWORD, graph=self.GRAPH, graphspace=self.GRAPHSPACE + url=self.URL, + user=self.USERNAME, + pwd=self.PASSWORD, + graph=self.GRAPH, + graphspace=self.GRAPHSPACE, ) assert self.client is not None From 6ad8fd9ba6ed488923ba7902dfdec7bbc064a53a Mon Sep 17 00:00:00 2001 From: Linyu <94553312+weijinglin@users.noreply.github.com> Date: Tue, 30 Sep 2025 10:55:47 +0800 Subject: [PATCH 59/71] refactor(RAG workflow): modularize flows, add streaming, and improve node initialization (#51) --- .../hugegraph_llm/demo/rag_demo/rag_block.py | 223 ++++++++++-------- .../src/hugegraph_llm/flows/common.py | 26 ++ .../flows/rag_flow_graph_only.py | 153 ++++++++++++ .../flows/rag_flow_graph_vector.py | 158 +++++++++++++ .../src/hugegraph_llm/flows/rag_flow_raw.py | 99 ++++++++ .../flows/rag_flow_vector_only.py | 123 ++++++++++ .../src/hugegraph_llm/flows/scheduler.py | 63 +++++ .../src/hugegraph_llm/nodes/base_node.py | 3 + .../nodes/common_node/merge_rerank_node.py | 83 +++++++ .../nodes/document_node/chunk_split.py | 2 +- .../hugegraph_node/commit_to_hugegraph.py | 3 +- .../nodes/hugegraph_node/fetch_graph_data.py | 3 +- .../nodes/hugegraph_node/graph_query_node.py | 93 ++++++++ .../nodes/hugegraph_node/schema.py | 3 +- .../nodes/index_node/build_semantic_index.py | 3 +- .../nodes/index_node/build_vector_index.py | 3 +- .../index_node/gremlin_example_index_query.py | 13 +- .../index_node/semantic_id_query_node.py | 91 +++++++ .../nodes/index_node/vector_query_node.py | 74 ++++++ .../nodes/llm_node/answer_synthesize_node.py | 99 ++++++++ .../nodes/llm_node/extract_info.py | 2 +- .../nodes/llm_node/keyword_extract_node.py | 80 +++++++ .../nodes/llm_node/prompt_generate.py | 2 +- .../nodes/llm_node/schema_build.py | 2 +- .../nodes/llm_node/text2gremlin.py | 10 +- .../src/hugegraph_llm/state/ai_state.py | 106 ++++++++- .../hugegraph_llm/utils/graph_index_utils.py | 115 ++------- 27 files changed, 1411 insertions(+), 224 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index 8f70c34bd..ca36867d9 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -21,12 +21,11 @@ from typing import AsyncGenerator, Literal, Optional, Tuple import gradio as gr +from hugegraph_llm.flows.scheduler import SchedulerSingleton import pandas as pd from gradio.utils import NamedString -from hugegraph_llm.config import huge_settings, llm_settings, prompt, resource_path -from hugegraph_llm.operators.graph_rag_task import RAGPipeline -from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize +from hugegraph_llm.config import resource_path, prompt, llm_settings from hugegraph_llm.utils.decorators import with_task_id from hugegraph_llm.utils.log import log @@ -72,44 +71,51 @@ def rag_answer( gr.Warning("Please select at least one generate mode.") return "", "", "", "" - rag = RAGPipeline() - if vector_search: - rag.query_vector_index() - if graph_search: - rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid( - vector_dis_threshold=vector_dis_threshold, - topk_per_keyword=topk_per_keyword, - ).import_schema(huge_settings.graph_name).query_graphdb( - num_gremlin_generate_example=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - max_graph_items=max_graph_items, - ) - # TODO: add more user-defined search strategies - rag.merge_dedup_rerank( - graph_ratio=graph_ratio, - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - topk_return_results=topk_return_results, - ) - rag.synthesize_answer( - raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt - ) - + scheduler = SchedulerSingleton.get_instance() try: - context = rag.run( - verbose=True, + # Select workflow by mode to avoid fetching the wrong pipeline from the pool + if graph_vector_answer or (graph_only_answer and vector_only_answer): + flow_key = "rag_graph_vector" + elif vector_only_answer: + flow_key = "rag_vector_only" + elif graph_only_answer: + flow_key = "rag_graph_only" + elif raw_answer: + flow_key = "rag_raw" + else: + raise RuntimeError("Unsupported flow type") + + res = scheduler.schedule_flow( + flow_key, query=text, vector_search=vector_search, graph_search=graph_search, + raw_answer=raw_answer, + vector_only_answer=vector_only_answer, + graph_only_answer=graph_only_answer, + graph_vector_answer=graph_vector_answer, + graph_ratio=graph_ratio, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + answer_prompt=answer_prompt, + keywords_extract_prompt=keywords_extract_prompt, + gremlin_tmpl_num=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, max_graph_items=max_graph_items, + topk_return_results=topk_return_results, + vector_dis_threshold=vector_dis_threshold, + topk_per_keyword=topk_per_keyword, ) - if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") + if res.get("switch_to_bleu"): + gr.Warning( + "Online reranker fails, automatically switches to local bleu rerank." + ) return ( - context.get("raw_answer", ""), - context.get("vector_only_answer", ""), - context.get("graph_only_answer", ""), - context.get("graph_vector_answer", ""), + res.get("raw_answer", ""), + res.get("vector_only_answer", ""), + res.get("graph_only_answer", ""), + res.get("graph_vector_answer", ""), ) except ValueError as e: log.critical(e) @@ -187,44 +193,47 @@ async def rag_answer_streaming( yield "", "", "", "" return - rag = RAGPipeline() - if vector_search: - rag.query_vector_index() - if graph_search: - rag.extract_keywords( - extract_template=keywords_extract_prompt - ).keywords_to_vid().import_schema(huge_settings.graph_name).query_graphdb( - num_gremlin_generate_example=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - ) - rag.merge_dedup_rerank( - graph_ratio, - rerank_method, - near_neighbor_first, - ) - # rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer, graph_vector_answer, answer_prompt) - try: - context = rag.run( - verbose=True, query=text, vector_search=vector_search, graph_search=graph_search - ) - if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") - answer_synthesize = AnswerSynthesize( + # Select the specific streaming workflow + scheduler = SchedulerSingleton.get_instance() + if graph_vector_answer or (graph_only_answer and vector_only_answer): + flow_key = "rag_graph_vector" + elif vector_only_answer: + flow_key = "rag_vector_only" + elif graph_only_answer: + flow_key = "rag_graph_only" + elif raw_answer: + flow_key = "rag_raw" + else: + raise RuntimeError("Unsupported flow type") + + async for res in scheduler.schedule_stream_flow( + flow_key, + query=text, + vector_search=vector_search, + graph_search=graph_search, raw_answer=raw_answer, vector_only_answer=vector_only_answer, graph_only_answer=graph_only_answer, graph_vector_answer=graph_vector_answer, - prompt_template=answer_prompt, - ) - async for context in answer_synthesize.run_streaming(context): - if context.get("switch_to_bleu"): - gr.Warning("Online reranker fails, automatically switches to local bleu rerank.") + graph_ratio=graph_ratio, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + answer_prompt=answer_prompt, + keywords_extract_prompt=keywords_extract_prompt, + gremlin_tmpl_num=gremlin_tmpl_num, + gremlin_prompt=gremlin_prompt, + ): + if res.get("switch_to_bleu"): + gr.Warning( + "Online reranker fails, automatically switches to local bleu rerank." + ) yield ( - context.get("raw_answer", ""), - context.get("vector_only_answer", ""), - context.get("graph_only_answer", ""), - context.get("graph_vector_answer", ""), + res.get("raw_answer", ""), + res.get("vector_only_answer", ""), + res.get("graph_only_answer", ""), + res.get("graph_vector_answer", ""), ) except ValueError as e: log.critical(e) @@ -242,7 +251,10 @@ def create_rag_block(): with gr.Column(scale=2): # with gr.Blocks().queue(max_size=20, default_concurrency_limit=5): inp = gr.Textbox( - value=prompt.default_question, label="Question", show_copy_button=True, lines=3 + value=prompt.default_question, + label="Question", + show_copy_button=True, + lines=3, ) # TODO: Only support inline formula now. Should support block formula @@ -271,7 +283,10 @@ def create_rag_block(): latex_delimiters=[{"left": "$", "right": "$", "display": False}], ) answer_prompt_input = gr.Textbox( - value=prompt.answer_prompt, label="Query Prompt", show_copy_button=True, lines=7 + value=prompt.answer_prompt, + label="Query Prompt", + show_copy_button=True, + lines=7, ) keywords_extract_prompt_input = gr.Textbox( value=prompt.keywords_extract_prompt, @@ -282,7 +297,9 @@ def create_rag_block(): with gr.Column(scale=1): with gr.Row(): - raw_radio = gr.Radio(choices=[True, False], value=False, label="Basic LLM Answer") + raw_radio = gr.Radio( + choices=[True, False], value=False, label="Basic LLM Answer" + ) vector_only_radio = gr.Radio( choices=[True, False], value=False, label="Vector-only Answer" ) @@ -306,7 +323,9 @@ def toggle_slider(enable): label="Rerank method", ) example_num = gr.Number( - value=-1, label="Template Num (<0 means disable text2gql) ", precision=0 + value=-1, + label="Template Num (<0 means disable text2gql) ", + precision=0, ) graph_ratio = gr.Slider( 0, 1, 0.6, label="Graph Ratio", step=0.1, interactive=False @@ -351,7 +370,7 @@ def toggle_slider(enable): """## 2. (Batch) Back-testing ) > 1. Download the template file & fill in the questions you want to test. > 2. Upload the file & click the button to generate answers. (Preview shows the first 40 lines) - > 3. The answer options are the same as the above RAG/Q&A frame + > 3. The answer options are the same as the above RAG/Q&A frame """ ) tests_df_headers = [ @@ -365,7 +384,9 @@ def toggle_slider(enable): # FIXME: "demo" might conflict with the graph name, it should be modified. answers_path = os.path.join(resource_path, "demo", "questions_answers.xlsx") questions_path = os.path.join(resource_path, "demo", "questions.xlsx") - questions_template_path = os.path.join(resource_path, "demo", "questions_template.xlsx") + questions_template_path = os.path.join( + resource_path, "demo", "questions_template.xlsx" + ) def read_file_to_excel(file: NamedString, line_count: Optional[int] = None): df = None @@ -412,20 +433,23 @@ def several_rag_answer( total_rows = len(df) for index, row in df.iterrows(): question = row.iloc[0] - basic_llm_answer, vector_only_answer, graph_only_answer, graph_vector_answer = ( - rag_answer( - question, - is_raw_answer, - is_vector_only_answer, - is_graph_only_answer, - is_graph_vector_answer, - graph_ratio_ui, - rerank_method_ui, - near_neighbor_first_ui, - custom_related_information_ui, - answer_prompt, - keywords_extract_prompt, - ) + ( + basic_llm_answer, + vector_only_answer, + graph_only_answer, + graph_vector_answer, + ) = rag_answer( + question, + is_raw_answer, + is_vector_only_answer, + is_graph_only_answer, + is_graph_vector_answer, + graph_ratio_ui, + rerank_method_ui, + near_neighbor_first_ui, + custom_related_information_ui, + answer_prompt, + keywords_extract_prompt, ) df.at[index, "Basic LLM Answer"] = basic_llm_answer df.at[index, "Vector-only Answer"] = vector_only_answer @@ -442,12 +466,18 @@ def several_rag_answer( file_types=[".xlsx", ".csv"], label="Questions File (.xlsx & csv)" ) with gr.Column(): - test_template_file = os.path.join(resource_path, "demo", "questions_template.xlsx") + test_template_file = os.path.join( + resource_path, "demo", "questions_template.xlsx" + ) gr.File(value=test_template_file, label="Download Template File") - answer_max_line_count = gr.Number(1, label="Max Lines To Show", minimum=1, maximum=40) + answer_max_line_count = gr.Number( + 1, label="Max Lines To Show", minimum=1, maximum=40 + ) answers_btn = gr.Button("Generate Answer (Batch)", variant="primary") # TODO: Set individual progress bars for dataframe - qa_dataframe = gr.DataFrame(label="Questions & Answers (Preview)", headers=tests_df_headers) + qa_dataframe = gr.DataFrame( + label="Questions & Answers (Preview)", headers=tests_df_headers + ) answers_btn.click( several_rag_answer, inputs=[ @@ -465,6 +495,15 @@ def several_rag_answer( ], outputs=[qa_dataframe, gr.File(label="Download Answered File", min_width=40)], ) - questions_file.change(read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count]) - answer_max_line_count.change(change_showing_excel, answer_max_line_count, qa_dataframe) - return inp, answer_prompt_input, keywords_extract_prompt_input, custom_related_information + questions_file.change( + read_file_to_excel, questions_file, [qa_dataframe, answer_max_line_count] + ) + answer_max_line_count.change( + change_showing_excel, answer_max_line_count, qa_dataframe + ) + return ( + inp, + answer_prompt_input, + keywords_extract_prompt_input, + custom_related_information, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/common.py b/hugegraph-llm/src/hugegraph_llm/flows/common.py index 4c552626a..e2348466c 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/common.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/common.py @@ -14,8 +14,10 @@ # limitations under the License. from abc import ABC, abstractmethod +from typing import Dict, Any, AsyncGenerator from hugegraph_llm.state.ai_state import WkFlowInput +from hugegraph_llm.utils.log import log class BaseFlow(ABC): @@ -43,3 +45,27 @@ def post_deal(self, *args, **kwargs): Post-processing interface. """ pass + + async def post_deal_stream( + self, pipeline=None + ) -> AsyncGenerator[Dict[str, Any], None]: + """ + Streaming post-processing interface. + Subclasses can override this method as needed. + """ + flow_name = self.__class__.__name__ + if pipeline is None: + yield {"error": "No pipeline provided"} + return + try: + state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info(f"{flow_name} post processing success") + stream_flow = state_json.get("stream_generator") + if stream_flow is None: + yield {"error": "No stream_generator found in workflow state"} + return + async for chunk in stream_flow: + yield chunk + except Exception as e: + log.error(f"{flow_name} post processing failed: {e}") + yield {"error": f"Post processing failed: {str(e)}"} diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py new file mode 100644 index 000000000..5feb3d471 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py @@ -0,0 +1,153 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +from typing import Optional, Literal + +from PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.llm_node.keyword_extract_node import KeywordExtractNode +from hugegraph_llm.nodes.index_node.semantic_id_query_node import SemanticIdQueryNode +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.nodes.hugegraph_node.graph_query_node import GraphQueryNode +from hugegraph_llm.nodes.common_node.merge_rerank_node import MergeRerankNode +from hugegraph_llm.nodes.llm_node.answer_synthesize_node import AnswerSynthesizeNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.utils.log import log + + +class RAGGraphOnlyFlow(BaseFlow): + """ + Workflow for graph-only answering (graph_only_answer) + """ + + def prepare( + self, + prepared_input: WkFlowInput, + query: str, + vector_search: bool = None, + graph_search: bool = None, + raw_answer: bool = None, + vector_only_answer: bool = None, + graph_only_answer: bool = None, + graph_vector_answer: bool = None, + graph_ratio: float = 0.5, + rerank_method: Literal["bleu", "reranker"] = "bleu", + near_neighbor_first: bool = False, + custom_related_information: str = "", + answer_prompt: Optional[str] = None, + keywords_extract_prompt: Optional[str] = None, + gremlin_tmpl_num: Optional[int] = -1, + gremlin_prompt: Optional[str] = None, + max_graph_items: int = None, + topk_return_results: int = None, + vector_dis_threshold: float = None, + topk_per_keyword: int = None, + **_: dict, + ): + prepared_input.query = query + prepared_input.vector_search = vector_search + prepared_input.graph_search = graph_search + prepared_input.raw_answer = raw_answer + prepared_input.vector_only_answer = vector_only_answer + prepared_input.graph_only_answer = graph_only_answer + prepared_input.graph_vector_answer = graph_vector_answer + prepared_input.gremlin_tmpl_num = gremlin_tmpl_num + prepared_input.gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt + prepared_input.max_graph_items = ( + max_graph_items or huge_settings.max_graph_items + ) + prepared_input.topk_per_keyword = ( + topk_per_keyword or huge_settings.topk_per_keyword + ) + prepared_input.topk_return_results = ( + topk_return_results or huge_settings.topk_return_results + ) + prepared_input.rerank_method = rerank_method + prepared_input.near_neighbor_first = near_neighbor_first + prepared_input.keywords_extract_prompt = ( + keywords_extract_prompt or prompt.keywords_extract_prompt + ) + prepared_input.answer_prompt = answer_prompt or prompt.answer_prompt + prepared_input.custom_related_information = custom_related_information + prepared_input.vector_dis_threshold = ( + vector_dis_threshold or huge_settings.vector_dis_threshold + ) + prepared_input.schema = huge_settings.graph_name + + prepared_input.data_json = { + "query": query, + "vector_search": vector_search, + "graph_search": graph_search, + "max_graph_items": max_graph_items or huge_settings.max_graph_items, + } + return + + def build_flow(self, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, **kwargs) + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + # Create nodes and register them with registerGElement + only_keyword_extract_node = KeywordExtractNode() + only_semantic_id_query_node = SemanticIdQueryNode() + only_schema_node = SchemaNode() + only_graph_query_node = GraphQueryNode() + merge_rerank_node = MergeRerankNode() + answer_synthesize_node = AnswerSynthesizeNode() + + pipeline.registerGElement(only_keyword_extract_node, set(), "only_keyword") + pipeline.registerGElement( + only_semantic_id_query_node, {only_keyword_extract_node}, "only_semantic" + ) + pipeline.registerGElement(only_schema_node, set(), "only_schema") + pipeline.registerGElement( + only_graph_query_node, + {only_schema_node, only_semantic_id_query_node}, + "only_graph", + ) + pipeline.registerGElement( + merge_rerank_node, {only_graph_query_node}, "merge_one" + ) + pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "graph") + log.info("RAGGraphOnlyFlow pipeline built successfully") + return pipeline + + def post_deal(self, pipeline=None): + if pipeline is None: + return json.dumps( + {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 + ) + try: + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGGraphOnlyFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } + except Exception as e: + log.error(f"RAGGraphOnlyFlow post processing failed: {e}") + return json.dumps( + {"error": f"Post processing failed: {str(e)}"}, + ensure_ascii=False, + indent=2, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py new file mode 100644 index 000000000..2f4a2bfa2 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py @@ -0,0 +1,158 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +from typing import Optional, Literal + +from PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.index_node.vector_query_node import VectorQueryNode +from hugegraph_llm.nodes.llm_node.keyword_extract_node import KeywordExtractNode +from hugegraph_llm.nodes.index_node.semantic_id_query_node import SemanticIdQueryNode +from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode +from hugegraph_llm.nodes.hugegraph_node.graph_query_node import GraphQueryNode +from hugegraph_llm.nodes.common_node.merge_rerank_node import MergeRerankNode +from hugegraph_llm.nodes.llm_node.answer_synthesize_node import AnswerSynthesizeNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.utils.log import log + + +class RAGGraphVectorFlow(BaseFlow): + """ + Workflow for graph + vector hybrid answering (graph_vector_answer) + """ + + def prepare( + self, + prepared_input: WkFlowInput, + query: str, + vector_search: bool = None, + graph_search: bool = None, + raw_answer: bool = None, + vector_only_answer: bool = None, + graph_only_answer: bool = None, + graph_vector_answer: bool = None, + graph_ratio: float = 0.5, + rerank_method: Literal["bleu", "reranker"] = "bleu", + near_neighbor_first: bool = False, + custom_related_information: str = "", + answer_prompt: Optional[str] = None, + keywords_extract_prompt: Optional[str] = None, + gremlin_tmpl_num: Optional[int] = -1, + gremlin_prompt: Optional[str] = None, + max_graph_items: int = None, + topk_return_results: int = None, + vector_dis_threshold: float = None, + topk_per_keyword: int = None, + **_: dict, + ): + prepared_input.query = query + prepared_input.vector_search = vector_search + prepared_input.graph_search = graph_search + prepared_input.raw_answer = raw_answer + prepared_input.vector_only_answer = vector_only_answer + prepared_input.graph_only_answer = graph_only_answer + prepared_input.graph_vector_answer = graph_vector_answer + prepared_input.graph_ratio = graph_ratio + prepared_input.gremlin_tmpl_num = gremlin_tmpl_num + prepared_input.gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt + prepared_input.max_graph_items = ( + max_graph_items or huge_settings.max_graph_items + ) + prepared_input.topk_return_results = ( + topk_return_results or huge_settings.topk_return_results + ) + prepared_input.topk_per_keyword = ( + topk_per_keyword or huge_settings.topk_per_keyword + ) + prepared_input.vector_dis_threshold = ( + vector_dis_threshold or huge_settings.vector_dis_threshold + ) + prepared_input.rerank_method = rerank_method + prepared_input.near_neighbor_first = near_neighbor_first + prepared_input.keywords_extract_prompt = ( + keywords_extract_prompt or prompt.keywords_extract_prompt + ) + prepared_input.answer_prompt = answer_prompt or prompt.answer_prompt + prepared_input.custom_related_information = custom_related_information + prepared_input.schema = huge_settings.graph_name + + prepared_input.data_json = { + "query": query, + "vector_search": vector_search, + "graph_search": graph_search, + "max_graph_items": max_graph_items or huge_settings.max_graph_items, + } + return + + def build_flow(self, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, **kwargs) + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + # Create nodes (registration style consistent with RAGFlow) + vector_query_node = VectorQueryNode() + keyword_extract_node = KeywordExtractNode() + semantic_id_query_node = SemanticIdQueryNode() + schema_node = SchemaNode() + graph_query_node = GraphQueryNode() + merge_rerank_node = MergeRerankNode() + answer_synthesize_node = AnswerSynthesizeNode() + + # Register nodes and their dependencies + pipeline.registerGElement(vector_query_node, set(), "vector") + pipeline.registerGElement(keyword_extract_node, set(), "keyword") + pipeline.registerGElement( + semantic_id_query_node, {keyword_extract_node}, "semantic" + ) + pipeline.registerGElement(schema_node, set(), "schema") + pipeline.registerGElement( + graph_query_node, {schema_node, semantic_id_query_node}, "graph" + ) + pipeline.registerGElement( + merge_rerank_node, {graph_query_node, vector_query_node}, "merge" + ) + pipeline.registerGElement( + answer_synthesize_node, {merge_rerank_node}, "graph_vector" + ) + log.info("RAGGraphVectorFlow pipeline built successfully") + return pipeline + + def post_deal(self, pipeline=None): + if pipeline is None: + return json.dumps( + {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 + ) + try: + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGGraphVectorFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } + except Exception as e: + log.error(f"RAGGraphVectorFlow post processing failed: {e}") + return json.dumps( + {"error": f"Post processing failed: {str(e)}"}, + ensure_ascii=False, + indent=2, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py new file mode 100644 index 000000000..f62e574bb --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py @@ -0,0 +1,99 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +from typing import Optional + +from PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.llm_node.answer_synthesize_node import AnswerSynthesizeNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.utils.log import log + + +class RAGRawFlow(BaseFlow): + """ + Workflow for basic LLM answering only (raw_answer) + """ + + def prepare( + self, + prepared_input: WkFlowInput, + query: str, + vector_search: bool = None, + graph_search: bool = None, + raw_answer: bool = None, + vector_only_answer: bool = None, + graph_only_answer: bool = None, + graph_vector_answer: bool = None, + custom_related_information: str = "", + answer_prompt: Optional[str] = None, + max_graph_items: int = None, + **_: dict, + ): + prepared_input.query = query + prepared_input.raw_answer = raw_answer + prepared_input.vector_only_answer = vector_only_answer + prepared_input.graph_only_answer = graph_only_answer + prepared_input.graph_vector_answer = graph_vector_answer + prepared_input.custom_related_information = custom_related_information + prepared_input.answer_prompt = answer_prompt or prompt.answer_prompt + prepared_input.schema = huge_settings.graph_name + + prepared_input.data_json = { + "query": query, + "vector_search": vector_search, + "graph_search": graph_search, + "max_graph_items": max_graph_items or huge_settings.max_graph_items, + } + return + + def build_flow(self, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, **kwargs) + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + # Create nodes and register with registerGElement (no GRegion required) + answer_synthesize_node = AnswerSynthesizeNode() + pipeline.registerGElement(answer_synthesize_node, set(), "raw") + log.info("RAGRawFlow pipeline built successfully") + return pipeline + + def post_deal(self, pipeline=None): + if pipeline is None: + return json.dumps( + {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 + ) + try: + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGRawFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } + except Exception as e: + log.error(f"RAGRawFlow post processing failed: {e}") + return json.dumps( + {"error": f"Post processing failed: {str(e)}"}, + ensure_ascii=False, + indent=2, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py new file mode 100644 index 000000000..c727eacce --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py @@ -0,0 +1,123 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +from typing import Optional, Literal + +from PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.nodes.index_node.vector_query_node import VectorQueryNode +from hugegraph_llm.nodes.common_node.merge_rerank_node import MergeRerankNode +from hugegraph_llm.nodes.llm_node.answer_synthesize_node import AnswerSynthesizeNode +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.utils.log import log + + +class RAGVectorOnlyFlow(BaseFlow): + """ + Workflow for vector-only answering (vector_only_answer) + """ + + def prepare( + self, + prepared_input: WkFlowInput, + query: str, + vector_search: bool = None, + graph_search: bool = None, + raw_answer: bool = None, + vector_only_answer: bool = None, + graph_only_answer: bool = None, + graph_vector_answer: bool = None, + rerank_method: Literal["bleu", "reranker"] = "bleu", + near_neighbor_first: bool = False, + custom_related_information: str = "", + answer_prompt: Optional[str] = None, + max_graph_items: int = None, + topk_return_results: int = None, + vector_dis_threshold: float = None, + **_: dict, + ): + prepared_input.query = query + prepared_input.vector_search = vector_search + prepared_input.graph_search = graph_search + prepared_input.raw_answer = raw_answer + prepared_input.vector_only_answer = vector_only_answer + prepared_input.graph_only_answer = graph_only_answer + prepared_input.graph_vector_answer = graph_vector_answer + prepared_input.vector_dis_threshold = ( + vector_dis_threshold or huge_settings.vector_dis_threshold + ) + prepared_input.topk_return_results = ( + topk_return_results or huge_settings.topk_return_results + ) + prepared_input.rerank_method = rerank_method + prepared_input.near_neighbor_first = near_neighbor_first + prepared_input.custom_related_information = custom_related_information + prepared_input.answer_prompt = answer_prompt or prompt.answer_prompt + prepared_input.schema = huge_settings.graph_name + + prepared_input.data_json = { + "query": query, + "vector_search": vector_search, + "graph_search": graph_search, + "max_graph_items": max_graph_items or huge_settings.max_graph_items, + } + return + + def build_flow(self, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, **kwargs) + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + # Create nodes (do not use GRegion, use registerGElement for all nodes) + only_vector_query_node = VectorQueryNode() + merge_rerank_node = MergeRerankNode() + answer_synthesize_node = AnswerSynthesizeNode() + + # Register nodes and dependencies, keep naming consistent with original + pipeline.registerGElement(only_vector_query_node, set(), "only_vector") + pipeline.registerGElement( + merge_rerank_node, {only_vector_query_node}, "merge_two" + ) + pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "vector") + log.info("RAGVectorOnlyFlow pipeline built successfully") + return pipeline + + def post_deal(self, pipeline=None): + if pipeline is None: + return json.dumps( + {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 + ) + try: + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGVectorOnlyFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } + except Exception as e: + log.error(f"RAGVectorOnlyFlow post processing failed: {e}") + return json.dumps( + {"error": f"Post processing failed: {str(e)}"}, + ensure_ascii=False, + indent=2, + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index 3aedbe7f2..5afa1bf8e 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -24,6 +24,11 @@ from hugegraph_llm.flows.get_graph_index_info import GetGraphIndexInfoFlow from hugegraph_llm.flows.build_schema import BuildSchemaFlow from hugegraph_llm.flows.prompt_generate import PromptGenerateFlow +from hugegraph_llm.flows.rag_flow_raw import RAGRawFlow +from hugegraph_llm.flows.rag_flow_vector_only import RAGVectorOnlyFlow +from hugegraph_llm.flows.rag_flow_graph_only import RAGGraphOnlyFlow +from hugegraph_llm.flows.rag_flow_graph_vector import RAGGraphVectorFlow +from hugegraph_llm.state.ai_state import WkFlowInput from hugegraph_llm.utils.log import log from hugegraph_llm.flows.text2gremlin import Text2GremlinFlow @@ -67,6 +72,23 @@ def __init__(self, max_pipeline: int = 10): "manager": GPipelineManager(), "flow": Text2GremlinFlow(), } + # New split rag pipelines + self.pipeline_pool["rag_raw"] = { + "manager": GPipelineManager(), + "flow": RAGRawFlow(), + } + self.pipeline_pool["rag_vector_only"] = { + "manager": GPipelineManager(), + "flow": RAGVectorOnlyFlow(), + } + self.pipeline_pool["rag_graph_only"] = { + "manager": GPipelineManager(), + "flow": RAGGraphOnlyFlow(), + } + self.pipeline_pool["rag_graph_vector"] = { + "manager": GPipelineManager(), + "flow": RAGGraphVectorFlow(), + } self.max_pipeline = max_pipeline # TODO: Implement Agentic Workflow @@ -108,6 +130,47 @@ def schedule_flow(self, flow: str, *args, **kwargs): manager.release(pipeline) return res + async def schedule_stream_flow(self, flow: str, *args, **kwargs): + if flow not in self.pipeline_pool: + raise ValueError(f"Unsupported workflow {flow}") + manager: GPipelineManager = self.pipeline_pool[flow]["manager"] + flow: BaseFlow = self.pipeline_pool[flow]["flow"] + pipeline: GPipeline = manager.fetch() + if pipeline is None: + # call coresponding flow_func to create new workflow + pipeline = flow.build_flow(*args, **kwargs) + try: + pipeline.getGParamWithNoEmpty("wkflow_input").stream = True + status = pipeline.init() + if status.isErr(): + error_msg = f"Error in flow init: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + status = pipeline.run() + if status.isErr(): + error_msg = f"Error in flow execution: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + async for res in flow.post_deal_stream(pipeline): + yield res + finally: + manager.add(pipeline) + else: + try: + # fetch pipeline & prepare input for flow + prepared_input: WkFlowInput = pipeline.getGParamWithNoEmpty( + "wkflow_input" + ) + prepared_input.stream = True + flow.prepare(prepared_input, *args, **kwargs) + status = pipeline.run() + if status.isErr(): + raise RuntimeError(f"Error in flow execution {status.getInfo()}") + async for res in flow.post_deal_stream(pipeline): + yield res + finally: + manager.release(pipeline) + class SchedulerSingleton: _instance = None diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py index 0ea0675c0..f90167305 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py @@ -30,6 +30,9 @@ def node_init(self): Node initialization method, can be overridden by subclasses. Returns a CStatus object indicating whether initialization succeeded. """ + if self.wk_input.data_json is not None: + self.context.assign_from_json(self.wk_input.data_json) + self.wk_input.data_json = None return CStatus() def run(self): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py new file mode 100644 index 000000000..78f53e231 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py @@ -0,0 +1,83 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict, Any +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.common_op.merge_dedup_rerank import MergeDedupRerank +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.utils.log import log + + +class MergeRerankNode(BaseNode): + """ + Merge and rerank node, responsible for merging vector and graph query results, deduplication and reranking. + """ + + operator: MergeDedupRerank + + def node_init(self): + """ + Initialize the merge and rerank operator. + """ + try: + # Read user configuration parameters from wk_input + embedding = get_embedding(llm_settings) + graph_ratio = self.wk_input.graph_ratio or 0.5 + rerank_method = self.wk_input.rerank_method or "bleu" + near_neighbor_first = self.wk_input.near_neighbor_first or False + custom_related_information = self.wk_input.custom_related_information or "" + topk_return_results = ( + self.wk_input.topk_return_results or huge_settings.topk_return_results + ) + + self.operator = MergeDedupRerank( + embedding=embedding, + graph_ratio=graph_ratio, + method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + topk_return_results=topk_return_results, + ) + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize MergeRerankNode: {e}") + from PyCGraph import CStatus + + return CStatus(-1, f"MergeRerankNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the merge and rerank operation. + """ + try: + # Perform merge, deduplication, and rerank + result = self.operator.run(data_json) + + # Log result statistics + vector_count = len(result.get("vector_result", [])) + graph_count = len(result.get("graph_result", [])) + merged_count = len(result.get("merged_result", [])) + + log.info( + f"Merge and rerank completed: {vector_count} vector results, " + f"{graph_count} graph results, {merged_count} merged results" + ) + + return result + + except Exception as e: + log.error(f"Merge and rerank failed: {e}") + return data_json 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 4c5acbe97..f71bd7bd5 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 @@ -37,7 +37,7 @@ def node_init(self): if isinstance(texts, str): texts = [texts] self.chunk_split_op = ChunkSplit(texts, split_type, language) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): 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 b576e8170..a4ebc7092 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 @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph import Commit2Graph from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -29,7 +28,7 @@ def node_init(self): if data_json: self.context.assign_from_json(data_json) self.commit_to_graph_op = Commit2Graph() - return CStatus() + return super().node_init() def operator_schedule(self, data_json): return self.commit_to_graph_op.run(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 b2434e524..99b428e5e 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 @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.hugegraph_op.fetch_graph_data import FetchGraphData from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -27,7 +26,7 @@ class FetchGraphDataNode(BaseNode): def node_init(self): self.fetch_graph_data_op = FetchGraphData(get_hg_client()) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): return self.fetch_graph_data_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py new file mode 100644 index 000000000..ae65ccb33 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py @@ -0,0 +1,93 @@ +# 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 PyCGraph import CStatus +from typing import Dict, Any +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery +from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.utils.log import log + + +class GraphQueryNode(BaseNode): + """ + Graph query node, responsible for retrieving relevant information from the graph database. + """ + + graph_rag_query: GraphRAGQuery + + def node_init(self): + """ + Initialize the graph query operator. + """ + try: + graph_name = huge_settings.graph_name + if not graph_name: + return CStatus(-1, "graph_name is required in wk_input") + + max_deep = self.wk_input.max_deep or 2 + max_graph_items = ( + self.wk_input.max_graph_items or huge_settings.max_graph_items + ) + max_v_prop_len = self.wk_input.max_v_prop_len or 2048 + max_e_prop_len = self.wk_input.max_e_prop_len or 256 + prop_to_match = self.wk_input.prop_to_match + num_gremlin_generate_example = self.wk_input.gremlin_tmpl_num or -1 + gremlin_prompt = ( + self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt + ) + + # Initialize GraphRAGQuery operator + self.graph_rag_query = GraphRAGQuery( + max_deep=max_deep, + max_graph_items=max_graph_items, + max_v_prop_len=max_v_prop_len, + max_e_prop_len=max_e_prop_len, + prop_to_match=prop_to_match, + num_gremlin_generate_example=num_gremlin_generate_example, + gremlin_prompt=gremlin_prompt, + ) + + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize GraphQueryNode: {e}") + + return CStatus(-1, f"GraphQueryNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the graph query operation. + """ + try: + # Get the query text from input + query = data_json.get("query", "") + + if not query: + log.warning("No query text provided for graph query") + return data_json + + # Execute the graph query (assuming schema and semantic query have been completed in previous nodes) + graph_result = self.graph_rag_query.run(data_json) + data_json.update(graph_result) + + log.info( + f"Graph query completed, found {len(data_json.get('graph_result', []))} results" + ) + + return data_json + + except Exception as e: + log.error(f"Graph query failed: {e}") + return data_json 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 84719d9eb..3face9d63 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py @@ -15,7 +15,6 @@ import json -from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.common_op.check_schema import CheckSchema from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager @@ -59,7 +58,7 @@ def node_init(self): else: log.info("Get schema '%s' from graphdb.", self.schema) self.schema_manager = self._import_schema(from_hugegraph=self.schema) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): log.debug("SchemaNode input state: %s", 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 ab31fa394..c01cffc91 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 @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from hugegraph_llm.config import llm_settings from hugegraph_llm.models.embeddings.init_embedding import get_embedding from hugegraph_llm.nodes.base_node import BaseNode @@ -28,7 +27,7 @@ class BuildSemanticIndexNode(BaseNode): def node_init(self): self.build_semantic_index_op = BuildSemanticIndex(get_embedding(llm_settings)) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): return self.build_semantic_index_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py index cf2f9b677..1f6a3c75b 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from hugegraph_llm.config import llm_settings from hugegraph_llm.models.embeddings.init_embedding import get_embedding from hugegraph_llm.nodes.base_node import BaseNode @@ -28,7 +27,7 @@ class BuildVectorIndexNode(BaseNode): def node_init(self): self.build_vector_index_op = BuildVectorIndex(get_embedding(llm_settings)) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): return self.build_vector_index_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py index eb033d869..e9283598a 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -19,9 +19,12 @@ from PyCGraph import CStatus +from hugegraph_llm.config import llm_settings from hugegraph_llm.nodes.base_node import BaseNode -from hugegraph_llm.operators.index_op.gremlin_example_index_query import GremlinExampleIndexQuery -from hugegraph_llm.models.embeddings.init_embedding import Embeddings +from hugegraph_llm.operators.index_op.gremlin_example_index_query import ( + GremlinExampleIndexQuery, +) +from hugegraph_llm.models.embeddings.init_embedding import get_embedding class GremlinExampleIndexQueryNode(BaseNode): @@ -29,13 +32,15 @@ class GremlinExampleIndexQueryNode(BaseNode): def node_init(self): # Build operator (index lazy-loading handled in operator) - embedding = Embeddings().get_embedding() + embedding = get_embedding(llm_settings) example_num = getattr(self.wk_input, "example_num", None) if not isinstance(example_num, int): example_num = 2 # Clamp to [0, 10] example_num = max(0, min(10, example_num)) - self.operator = GremlinExampleIndexQuery(embedding=embedding, num_examples=example_num) + self.operator = GremlinExampleIndexQuery( + embedding=embedding, num_examples=example_num + ) return CStatus() def operator_schedule(self, data_json: Dict[str, Any]): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py new file mode 100644 index 000000000..bf605aa49 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -0,0 +1,91 @@ +# 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 PyCGraph import CStatus +from typing import Dict, Any +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.semantic_id_query import SemanticIdQuery +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.utils.log import log + + +class SemanticIdQueryNode(BaseNode): + """ + Semantic ID query node, responsible for semantic matching based on keywords. + """ + + semantic_id_query: SemanticIdQuery + + def node_init(self): + """ + Initialize the semantic ID query operator. + """ + try: + graph_name = huge_settings.graph_name + if not graph_name: + return CStatus(-1, "graph_name is required in wk_input") + + embedding = get_embedding(llm_settings) + by = self.wk_input.semantic_by or "keywords" + topk_per_keyword = ( + self.wk_input.topk_per_keyword or huge_settings.topk_per_keyword + ) + topk_per_query = self.wk_input.topk_per_query or 10 + vector_dis_threshold = ( + self.wk_input.vector_dis_threshold or huge_settings.vector_dis_threshold + ) + + # Initialize the semantic ID query operator + self.semantic_id_query = SemanticIdQuery( + embedding=embedding, + by=by, + topk_per_keyword=topk_per_keyword, + topk_per_query=topk_per_query, + vector_dis_threshold=vector_dis_threshold, + ) + + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize SemanticIdQueryNode: {e}") + + return CStatus(-1, f"SemanticIdQueryNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the semantic ID query operation. + """ + try: + # Get the query text and keywords from input + query = data_json.get("query", "") + keywords = data_json.get("keywords", []) + + if not query and not keywords: + log.warning("No query text or keywords provided for semantic query") + return data_json + + # Perform the semantic query + semantic_result = self.semantic_id_query.run(data_json) + + match_vids = semantic_result.get("match_vids", []) + log.info( + f"Semantic query completed, found {len(match_vids)} matching vertex IDs" + ) + + return semantic_result + + except Exception as e: + log.error(f"Semantic query failed: {e}") + return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py new file mode 100644 index 000000000..48b50acf3 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.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. + +from typing import Dict, Any +from hugegraph_llm.config import llm_settings +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.utils.log import log + + +class VectorQueryNode(BaseNode): + """ + Vector query node, responsible for retrieving relevant documents from the vector index + """ + + operator: VectorIndexQuery + + def node_init(self): + """ + Initialize the vector query operator + """ + try: + # 从 wk_input 中读取用户配置参数 + embedding = get_embedding(llm_settings) + max_items = ( + self.wk_input.max_items if self.wk_input.max_items is not None else 3 + ) + + self.operator = VectorIndexQuery(embedding=embedding, topk=max_items) + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize VectorQueryNode: {e}") + from PyCGraph import CStatus + + return CStatus(-1, f"VectorQueryNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the vector query operation + """ + try: + # Get the query text from input + query = data_json.get("query", "") + if not query: + log.warning("No query text provided for vector query") + return data_json + + # Perform the vector query + result = self.operator.run({"query": query}) + + # Update the state + data_json.update(result) + log.info( + f"Vector query completed, found {len(result.get('vector_result', []))} results" + ) + + return data_json + + except Exception as e: + log.error(f"Vector query failed: {e}") + return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py new file mode 100644 index 000000000..22b970b4a --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py @@ -0,0 +1,99 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict, Any +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize +from hugegraph_llm.utils.log import log + + +class AnswerSynthesizeNode(BaseNode): + """ + Answer synthesis node, responsible for generating the final answer based on retrieval results. + """ + + operator: AnswerSynthesize + + def node_init(self): + """ + Initialize the answer synthesis operator. + """ + try: + prompt_template = self.wk_input.answer_prompt + raw_answer = self.wk_input.raw_answer or False + vector_only_answer = self.wk_input.vector_only_answer or False + graph_only_answer = self.wk_input.graph_only_answer or False + graph_vector_answer = self.wk_input.graph_vector_answer or False + + self.operator = AnswerSynthesize( + prompt_template=prompt_template, + raw_answer=raw_answer, + vector_only_answer=vector_only_answer, + graph_only_answer=graph_only_answer, + graph_vector_answer=graph_vector_answer, + ) + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize AnswerSynthesizeNode: {e}") + from PyCGraph import CStatus + + return CStatus(-1, f"AnswerSynthesizeNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the answer synthesis operation. + """ + try: + if self.getGParamWithNoEmpty("wkflow_input").stream: + # Streaming mode: return a generator for streaming output + data_json["stream_generator"] = self.operator.run_streaming(data_json) + return data_json + else: + # Non-streaming mode: execute answer synthesis + result = self.operator.run(data_json) + + # Record the types of answers generated + answer_types = [] + if result.get("raw_answer"): + answer_types.append("raw") + if result.get("vector_only_answer"): + answer_types.append("vector_only") + if result.get("graph_only_answer"): + answer_types.append("graph_only") + if result.get("graph_vector_answer"): + answer_types.append("graph_vector") + + log.info( + f"Answer synthesis completed for types: {', '.join(answer_types)}" + ) + + # Print enabled answer types according to self.wk_input configuration + wk_input_types = [] + if getattr(self.wk_input, "raw_answer", False): + wk_input_types.append("raw") + if getattr(self.wk_input, "vector_only_answer", False): + wk_input_types.append("vector_only") + if getattr(self.wk_input, "graph_only_answer", False): + wk_input_types.append("graph_only") + if getattr(self.wk_input, "graph_vector_answer", False): + wk_input_types.append("graph_vector") + log.info( + f"Enabled answer types according to wk_input config: {', '.join(wk_input_types)}" + ) + return result + + except Exception as e: + log.error(f"Answer synthesis failed: {e}") + return 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 8bceed804..628765f58 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 @@ -43,7 +43,7 @@ def node_init(self): self.property_graph_extract = PropertyGraphExtract(llm, example_prompt) else: return CStatus(-1, f"Unsupported extract_type: {extract_type}") - return CStatus() + return super().node_init() def operator_schedule(self, data_json): if self.extract_type == "triples": diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py new file mode 100644 index 000000000..76fc06eb3 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py @@ -0,0 +1,80 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict, Any +from PyCGraph import CStatus + +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.llm_op.keyword_extract import KeywordExtract +from hugegraph_llm.utils.log import log + + +class KeywordExtractNode(BaseNode): + operator: KeywordExtract + + """ + Keyword extraction node, responsible for extracting keywords from query text. + """ + + def node_init(self): + """ + Initialize the keyword extraction operator. + """ + try: + max_keywords = ( + self.wk_input.max_keywords + if self.wk_input.max_keywords is not None + else 5 + ) + language = ( + self.wk_input.language + if self.wk_input.language is not None + else "english" + ) + extract_template = self.wk_input.keywords_extract_prompt + + self.operator = KeywordExtract( + text=self.wk_input.query, + max_keywords=max_keywords, + language=language, + extract_template=extract_template, + ) + return super().node_init() + except Exception as e: + log.error(f"Failed to initialize KeywordExtractNode: {e}") + return CStatus(-1, f"KeywordExtractNode initialization failed: {e}") + + def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Execute the keyword extraction operation. + """ + try: + # Perform keyword extraction + result = self.operator.run(data_json) + if "keywords" not in result: + log.warning("Keyword extraction result missing 'keywords' field") + result["keywords"] = [] + + log.info(f"Extracted keywords: {result.get('keywords', [])}") + + return result + + except Exception as e: + log.error(f"Keyword extraction failed: {e}") + # Add error flag to indicate failure + error_result = data_json.copy() + error_result["error"] = str(e) + error_result["keywords"] = [] + return error_result diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py index 317f9e6ac..8c49994fd 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/prompt_generate.py @@ -50,7 +50,7 @@ def node_init(self): "example_name": self.wk_input.example_name, } self.context.assign_from_json(context) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): """ diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py index 7df2e68e7..408adb10a 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py @@ -75,7 +75,7 @@ def node_init(self): } self.context.assign_from_json(_context_payload) - return CStatus() + return super().node_init() def operator_schedule(self, data_json): try: diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py index ffbafbaf4..a36831526 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py @@ -22,13 +22,15 @@ from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.config import prompt as prompt_cfg +from hugegraph_llm.models.llms.init_llm import get_text2gql_llm +from hugegraph_llm.config import llm_settings, prompt as prompt_cfg def _stable_schema_string(state_json: Dict[str, Any]) -> str: if "simple_schema" in state_json and state_json["simple_schema"] is not None: - return json.dumps(state_json["simple_schema"], ensure_ascii=False, sort_keys=True) + return json.dumps( + state_json["simple_schema"], ensure_ascii=False, sort_keys=True + ) if "schema" in state_json and state_json["schema"] is not None: return json.dumps(state_json["schema"], ensure_ascii=False, sort_keys=True) return "" @@ -39,7 +41,7 @@ class Text2GremlinNode(BaseNode): def node_init(self): # Select LLM - llm = LLMs().get_text2gql_llm() + llm = get_text2gql_llm(llm_settings) # Serialize schema deterministically state_json = self.context.to_json() schema_str = _stable_schema_string(state_json) diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index f941098b1..3a6fd3c1c 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -24,7 +24,6 @@ class WkFlowInput(GParam): split_type: str = None # split type used by ChunkSplit Node example_prompt: str = None # need by graph information extract schema: str = None # Schema information requeired by SchemaNode - graph_name: str = None data_json = None extract_type = None query_examples = None @@ -34,11 +33,45 @@ class WkFlowInput(GParam): scenario: str = None # Scenario description example_name: str = None # Example name # Fields for Text2Gremlin - query: str = None example_num: int = None gremlin_prompt: str = None requested_outputs: Optional[List[str]] = None + # RAG Flow related fields + query: str = None # User query for RAG + vector_search: bool = None # Enable vector search + graph_search: bool = None # Enable graph search + raw_answer: bool = None # Return raw answer + vector_only_answer: bool = None # Vector only answer mode + graph_only_answer: bool = None # Graph only answer mode + graph_vector_answer: bool = None # Combined graph and vector answer + graph_ratio: float = None # Graph ratio for merging + rerank_method: str = None # Reranking method + near_neighbor_first: bool = None # Near neighbor first flag + custom_related_information: str = None # Custom related information + answer_prompt: str = None # Answer generation prompt + keywords_extract_prompt: str = None # Keywords extraction prompt + gremlin_tmpl_num: int = None # Gremlin template number + gremlin_prompt: str = None # Gremlin generation prompt + max_graph_items: int = None # Maximum graph items + topk_return_results: int = None # Top-k return results + vector_dis_threshold: float = None # Vector distance threshold + topk_per_keyword: int = None # Top-k per keyword + max_keywords: int = None + max_items: int = None + + # Semantic query related fields + semantic_by: str = None # Semantic query method + topk_per_query: int = None # Top-k per query + + # Graph query related fields + max_deep: int = None # Maximum depth for graph traversal + max_v_prop_len: int = None # Maximum vertex property length + max_e_prop_len: int = None # Maximum edge property length + prop_to_match: str = None # Property to match + + stream: bool = None # used for recognize stream mode + def reset(self, _: CStatus) -> None: self.texts = None self.language = None @@ -55,10 +88,40 @@ def reset(self, _: CStatus) -> None: self.scenario = None self.example_name = None # Text2Gremlin related configuration - self.query = None self.example_num = None self.gremlin_prompt = None self.requested_outputs = None + # RAG Flow related fields + self.query = None + self.vector_search = None + self.graph_search = None + self.raw_answer = None + self.vector_only_answer = None + self.graph_only_answer = None + self.graph_vector_answer = None + self.graph_ratio = None + self.rerank_method = None + self.near_neighbor_first = None + self.custom_related_information = None + self.answer_prompt = None + self.keywords_extract_prompt = None + self.gremlin_tmpl_num = None + self.gremlin_prompt = None + self.max_graph_items = None + self.topk_return_results = None + self.vector_dis_threshold = None + self.topk_per_keyword = None + self.max_keywords = None + self.max_items = None + # Semantic query related fields + self.semantic_by = None + self.topk_per_query = None + # Graph query related fields + self.max_deep = None + self.max_v_prop_len = None + self.max_e_prop_len = None + self.prop_to_match = None + self.stream = None class WkFlowState(GParam): @@ -83,6 +146,17 @@ class WkFlowState(GParam): template_exec_res: Optional[Any] = None raw_exec_res: Optional[Any] = None + match_vids = None + vector_result = None + graph_result = None + + raw_answer: str = None + vector_only_answer: str = None + graph_only_answer: str = None + graph_vector_answer: str = None + + merged_result = None + def setup(self): self.schema = None self.simple_schema = None @@ -90,7 +164,7 @@ def setup(self): self.edges = None self.vertices = None self.triples = None - self.call_count = 0 + self.call_count = None self.keywords = None self.vector_result = None @@ -99,12 +173,20 @@ def setup(self): self.generated_extract_prompt = None # Text2Gremlin results reset - self.match_result = [] - self.result = "" - self.raw_result = "" - self.template_exec_res = "" - self.raw_exec_res = "" + self.match_result = None + self.result = None + self.raw_result = None + self.template_exec_res = None + self.raw_exec_res = None + self.raw_answer = None + self.vector_only_answer = None + self.graph_only_answer = None + self.graph_vector_answer = None + + self.vector_result = None + self.graph_result = None + self.merged_result = None return CStatus() def to_json(self): @@ -116,7 +198,11 @@ def to_json(self): dict: A dictionary containing non-None instance members and their serialized values. """ # Only export instance attributes (excluding methods and class attributes) whose values are not None - return {k: v for k, v in self.__dict__.items() if not k.startswith("_") and v is not None} + return { + k: v + for k, v in self.__dict__.items() + if not k.startswith("_") and v is not None + } # Implement a method that assigns keys from data_json as WkFlowState member variables def assign_from_json(self, data_json: dict): diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 7b870033a..3f527f2fa 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -44,30 +44,17 @@ def get_graph_index_info(): raise gr.Error(str(e)) -def get_graph_index_info_old(): - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - graph_summary_info = builder.fetch_graph_data().run() - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) - index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, getattr(builder.embedding, "model_name", None) - ) - vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) - graph_summary_info["vid_index"] = { - "embed_dim": vector_index.index.d, - "num_vectors": vector_index.index.ntotal, - "num_vids": len(vector_index.properties), - } - return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) - - def clean_all_graph_index(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(Embeddings().get_embedding(), "model_name", None), ) - VectorIndex.clean(str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix) + VectorIndex.clean( + str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix + ) VectorIndex.clean( str(os.path.join(resource_path, folder_name, "gremlin_examples")), filename_prefix, @@ -99,14 +86,18 @@ def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: def extract_graph_origin(input_file, input_text, schema, example_prompt) -> str: texts = read_documents(input_file, input_text) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) + builder = KgBuilder( + LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() + ) if not schema: return "ERROR: please input with correct schema/format." error_message = parse_schema(schema, builder) if error_message: return error_message - builder.chunk_split(texts, "document", "zh").extract_info(example_prompt, "property_graph") + builder.chunk_split(texts, "document", "zh").extract_info( + example_prompt, "property_graph" + ) try: context = builder.run() @@ -155,20 +146,6 @@ def update_vid_embedding(): raise gr.Error(str(e)) -def update_vid_embedding_old(): - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - builder.fetch_graph_data().build_vertex_id_semantic_index() - log.debug("Operators: %s", builder.operators) - try: - context = builder.run() - removed_num = context["removed_vid_vector_num"] - added_num = context["added_vid_vector_num"] - return f"Removed {removed_num} vectors, added {added_num} vectors." - except Exception as e: # pylint: disable=broad-exception-caught - log.error(e) - raise gr.Error(str(e)) - - def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: scheduler = SchedulerSingleton.get_instance() @@ -181,73 +158,11 @@ def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: return data -def import_graph_data_old(data: str, schema: str) -> Union[str, Dict[str, Any]]: - try: - data_json = json.loads(data.strip()) - log.debug("Import graph data: %s", data) - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - if schema: - error_message = parse_schema(schema, builder) - if error_message: - return error_message - - context = builder.commit_to_hugegraph().run(data_json) - gr.Info("Import graph data successfully!") - print(context) - return json.dumps(context, ensure_ascii=False, indent=2) - except Exception as e: # pylint: disable=W0718 - log.error(e) - traceback.print_exc() - # Note: can't use gr.Error here - gr.Warning(str(e) + " Please check the graph data format/type carefully.") - return data - - def build_schema(input_text, query_example, few_shot): scheduler = SchedulerSingleton.get_instance() try: - return scheduler.schedule_flow("build_schema", input_text, query_example, few_shot) + return scheduler.schedule_flow( + "build_schema", input_text, query_example, few_shot + ) except (TypeError, ValueError) as e: raise gr.Error(f"Schema generation failed: {e}") - - -def build_schema_old(input_text, query_example, few_shot): - context = { - "raw_texts": [input_text] if input_text else [], - "query_examples": [], - "few_shot_schema": {}, - } - - if few_shot: - try: - context["few_shot_schema"] = json.loads(few_shot) - except json.JSONDecodeError as e: - raise gr.Error(f"Few Shot Schema is not in a valid JSON format: {e}") from e - - if query_example: - try: - parsed_examples = json.loads(query_example) - # Validate and retain the description and gremlin fields - context["query_examples"] = [ - { - "description": ex.get("description", ""), - "gremlin": ex.get("gremlin", ""), - } - for ex in parsed_examples - if isinstance(ex, dict) and "description" in ex and "gremlin" in ex - ] - except json.JSONDecodeError as e: - raise gr.Error(f"Query Examples is not in a valid JSON format: {e}") from e - - builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client()) - try: - schema = builder.build_schema().run(context) - except Exception as e: - log.error("Failed to generate schema: %s", e) - raise gr.Error(f"Schema generation failed: {e}") from e - try: - formatted_schema = json.dumps(schema, ensure_ascii=False, indent=2) - return formatted_schema - except (TypeError, ValueError) as e: - log.error("Failed to format schema: %s", e) - return str(schema) From edae38950adb51b47352b6376d4b569f99f2d1e9 Mon Sep 17 00:00:00 2001 From: jinglinwei Date: Mon, 13 Oct 2025 01:01:47 +0800 Subject: [PATCH 60/71] refactor: port batch build gremlin examples & delete some doc related to Pipeline(old design) & refactor some operator's design and implementation & code format --- README.md | 36 -- hugegraph-llm/README.md | 155 +++--- hugegraph-llm/pyproject.toml | 6 +- .../src/hugegraph_llm/api/admin_api.py | 2 +- .../src/hugegraph_llm/api/rag_api.py | 72 ++- .../hugegraph_llm/demo/rag_demo/rag_block.py | 36 +- .../demo/rag_demo/text2gremlin_block.py | 158 ++---- .../demo/rag_demo/vector_graph_block.py | 59 ++- .../src/hugegraph_llm/flows/__init__.py | 18 + .../flows/build_example_index.py | 62 +++ .../src/hugegraph_llm/flows/build_schema.py | 16 +- .../hugegraph_llm/flows/build_vector_index.py | 16 +- .../src/hugegraph_llm/flows/common.py | 29 +- .../flows/get_graph_index_info.py | 16 +- .../src/hugegraph_llm/flows/graph_extract.py | 25 +- .../hugegraph_llm/flows/import_graph_data.py | 8 +- .../hugegraph_llm/flows/prompt_generate.py | 19 +- .../flows/rag_flow_graph_only.py | 110 ++-- .../flows/rag_flow_graph_vector.py | 55 +- .../src/hugegraph_llm/flows/rag_flow_raw.py | 49 +- .../flows/rag_flow_vector_only.py | 53 +- .../src/hugegraph_llm/flows/scheduler.py | 118 +++-- .../src/hugegraph_llm/flows/text2gremlin.py | 14 +- .../flows/update_vid_embeddings.py | 18 +- .../src/hugegraph_llm/flows/utils.py | 34 -- .../models/embeddings/init_embedding.py | 28 +- .../src/hugegraph_llm/models/llms/init_llm.py | 96 ++-- .../src/hugegraph_llm/nodes/base_node.py | 43 +- .../nodes/common_node/merge_rerank_node.py | 14 +- .../nodes/document_node/chunk_split.py | 2 +- .../nodes/hugegraph_node/fetch_graph_data.py | 9 +- .../nodes/hugegraph_node/graph_query_node.py | 486 ++++++++++++++++-- .../nodes/hugegraph_node/schema.py | 26 +- .../index_node/build_gremlin_example_index.py | 43 ++ .../index_node/semantic_id_query_node.py | 92 ++-- .../nodes/index_node/vector_query_node.py | 27 +- .../nodes/llm_node/answer_synthesize_node.py | 103 ++-- .../nodes/llm_node/extract_info.py | 3 +- .../nodes/llm_node/keyword_extract_node.py | 39 +- .../nodes/llm_node/schema_build.py | 10 +- .../nodes/llm_node/text2gremlin.py | 18 +- hugegraph-llm/src/hugegraph_llm/nodes/util.py | 29 +- .../operators/gremlin_generate_task.py | 81 --- .../hugegraph_op/commit_to_hugegraph.py | 59 ++- .../operators/hugegraph_op/graph_rag_query.py | 455 ---------------- .../index_op/build_semantic_index.py | 24 +- .../operators/kg_construction_task.py | 120 ----- .../operators/llm_op/keyword_extract.py | 58 +-- .../{graph_rag_task.py => operator_list.py} | 233 +++++---- .../src/hugegraph_llm/operators/util.py | 27 - .../src/hugegraph_llm/state/ai_state.py | 186 ++++--- .../hugegraph_llm/utils/graph_index_utils.py | 96 ++-- .../hugegraph_llm/utils/vector_index_utils.py | 19 +- hugegraph-ml/pyproject.toml | 2 +- hugegraph-python-client/pyproject.toml | 2 +- pyproject.toml | 2 +- scripts/build_llm_image.sh | 2 +- .../hugegraph-llm/fixed_flow/design.md | 3 +- .../hugegraph-llm/fixed_flow/requirements.md | 0 .../hugegraph-llm/fixed_flow/tasks.md | 0 style/pylint.conf | 4 +- vermeer-python-client/pyproject.toml | 4 +- 62 files changed, 1760 insertions(+), 1869 deletions(-) create mode 100644 hugegraph-llm/src/hugegraph_llm/flows/build_example_index.py delete mode 100644 hugegraph-llm/src/hugegraph_llm/flows/utils.py create mode 100644 hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py delete mode 100644 hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py delete mode 100644 hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py delete mode 100644 hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py rename hugegraph-llm/src/hugegraph_llm/operators/{graph_rag_task.py => operator_list.py} (54%) delete mode 100644 hugegraph-llm/src/hugegraph_llm/operators/util.py mode change 100644 => 100755 scripts/build_llm_image.sh rename {.vibedev/spec => spec}/hugegraph-llm/fixed_flow/design.md (99%) rename {.vibedev/spec => spec}/hugegraph-llm/fixed_flow/requirements.md (100%) rename {.vibedev/spec => spec}/hugegraph-llm/fixed_flow/tasks.md (100%) diff --git a/README.md b/README.md index 14f02ca1c..a495968ec 100644 --- a/README.md +++ b/README.md @@ -75,42 +75,6 @@ python -m hugegraph_llm.demo.rag_demo.app > [!NOTE] > Examples assume you've activated the virtual environment with `source .venv/bin/activate` -#### GraphRAG - Question Answering - -```python -from hugegraph_llm.operators.graph_rag_task import RAGPipeline - -# Initialize RAG pipeline -graph_rag = RAGPipeline() - -# Ask questions about your graph -result = (graph_rag - .extract_keywords(text="Tell me about Al Pacino.") - .keywords_to_vid() - .query_graphdb(max_deep=2, max_graph_items=30) - .merge_dedup_rerank() - .synthesize_answer() - .run()) -``` - -#### Knowledge Graph Construction - -```python -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.kg_construction_task import KgBuilder - -# Build KG from text -TEXT = "Your text content here..." -builder = KgBuilder(LLMs().get_chat_llm()) - -(builder - .import_schema(from_hugegraph="hugegraph") - .chunk_split(TEXT) - .extract_info(extract_type="property_graph") - .commit_to_hugegraph() - .run()) -``` - #### Graph Machine Learning ```bash diff --git a/hugegraph-llm/README.md b/hugegraph-llm/README.md index 526320d4a..8b7e15c50 100644 --- a/hugegraph-llm/README.md +++ b/hugegraph-llm/README.md @@ -89,14 +89,14 @@ curl -LsSf https://astral.sh/uv/install.sh | sh # 3. Clone and setup project git clone https://github.com/apache/incubator-hugegraph-ai.git -cd incubator-hugegraph-ai/hugegraph-llm +cd incubator-hugegraph-ai # Configure environment (see config.md for detailed options), .env will auto create if not exists # 4. Install dependencies and activate environment # NOTE: If download is slow, uncomment mirror lines in ../pyproject.toml or use: uv config --global index.url https://pypi.tuna.tsinghua.edu.cn/simple # Or create local uv.toml with mirror settings to avoid git diff (see uv.toml example in root) -uv sync # Automatically creates .venv and installs dependencies +uv sync --extra llm # Automatically creates .venv and installs dependencies source .venv/bin/activate # Activate once - all commands below assume this environment # 5. Launch RAG demo @@ -146,84 +146,6 @@ Use the Gradio interface for visual knowledge graph building: ![Knowledge Graph Builder](https://hugegraph.apache.org/docs/images/gradio-kg.png) -#### Programmatic Construction - -Build knowledge graphs with code using the `KgBuilder` class: - -```python -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.kg_construction_task import KgBuilder - -# Initialize and chain operations -TEXT = "Your input text here..." -builder = KgBuilder(LLMs().get_chat_llm()) - -( - builder - .import_schema(from_hugegraph="talent_graph").print_result() - .chunk_split(TEXT).print_result() - .extract_info(extract_type="property_graph").print_result() - .commit_to_hugegraph() - .run() -) -``` - -**Pipeline Workflow:** - -```mermaid -graph LR - A[Import Schema] --> B[Chunk Split] - B --> C[Extract Info] - C --> D[Commit to HugeGraph] - D --> E[Execute Pipeline] - - style A fill:#fff2cc - style B fill:#d5e8d4 - style C fill:#dae8fc - style D fill:#f8cecc - style E fill:#e1d5e7 -``` - -### Graph-Enhanced RAG - -Leverage HugeGraph for retrieval-augmented generation: - -```python -from hugegraph_llm.operators.graph_rag_task import RAGPipeline - -# Initialize RAG pipeline -graph_rag = RAGPipeline() - -# Execute RAG workflow -( - graph_rag - .extract_keywords(text="Tell me about Al Pacino.") - .keywords_to_vid() - .query_graphdb(max_deep=2, max_graph_items=30) - .merge_dedup_rerank() - .synthesize_answer(vector_only_answer=False, graph_only_answer=True) - .run(verbose=True) -) -``` - -**RAG Pipeline Flow:** - -```mermaid -graph TD - A[User Query] --> B[Extract Keywords] - B --> C[Match Graph Nodes] - C --> D[Retrieve Graph Context] - D --> E[Rerank Results] - E --> F[Generate Answer] - - style A fill:#e3f2fd - style B fill:#f3e5f5 - style C fill:#e8f5e8 - style D fill:#fff3e0 - style E fill:#fce4ec - style F fill:#e0f2f1 -``` - ## 🔧 Configuration After running the demo, configuration files are automatically generated: @@ -248,6 +170,79 @@ The system supports both English and Chinese prompts. To switch languages: **LLM Provider Support**: This project uses [LiteLLM](https://docs.litellm.ai/docs/providers) for multi-provider LLM support. +### Programmatic Examples (new workflow engine) + +If you previously used high-level classes like `RAGPipeline` or `KgBuilder`, the project now exposes stable flows through the `Scheduler` API. Use `SchedulerSingleton.get_instance().schedule_flow(...)` to invoke workflows programmatically. Below are concise, working examples that match the new architecture. + +1) RAG (graph-only) query example + +```python +from hugegraph_llm.flows.scheduler import SchedulerSingleton + +scheduler = SchedulerSingleton.get_instance() +res = scheduler.schedule_flow( + "rag_graph_only", + query="Tell me about Al Pacino.", + graph_only_answer=True, + vector_only_answer=False, + raw_answer=False, + gremlin_tmpl_num=-1, + gremlin_prompt=None, +) + +print(res.get("graph_only_answer")) +``` + +2) RAG (vector-only) query example + +```python +from hugegraph_llm.flows.scheduler import SchedulerSingleton + +scheduler = SchedulerSingleton.get_instance() +res = scheduler.schedule_flow( + "rag_vector_only", + query="Summarize the career of Ada Lovelace.", + vector_only_answer=True, + vector_search=True +) + +print(res.get("vector_only_answer")) +``` + +3) Text -> Gremlin (text2gremlin) example + +```python +from hugegraph_llm.flows.scheduler import SchedulerSingleton + +scheduler = SchedulerSingleton.get_instance() +response = scheduler.schedule_flow( + "text2gremlin", + "find people who worked with Alan Turing", + 2, # example_num + "hugegraph", # schema_input (graph name or schema) + None, # gremlin_prompt_input (optional) + ["template_gremlin", "raw_gremlin"], +) + +print(response.get("template_gremlin")) +``` + +4) Build example index (used by text2gremlin examples) + +```python +from hugegraph_llm.flows.scheduler import SchedulerSingleton + +examples = [{"id": "natural language query", "gremlin": "g.V().hasLabel('person').valueMap()"}] +res = SchedulerSingleton.get_instance().schedule_flow("build_examples_index", examples) +print(res) +``` + +### Migration guide: RAGPipeline / KgBuilder → Scheduler flows + +Why the change: the internal execution engine was refactored to a pipeline-based scheduler (GPipeline + GPipelineManager). The scheduler provides a stable entrypoint while keeping flow implementations modular. + +If you need help migrating a specific snippet, open a PR or issue and include the old code — we can provide a targeted conversion. + ## 🤖 Developer Guidelines > [!IMPORTANT] > **For developers contributing to hugegraph-llm with AI coding assistance:** diff --git a/hugegraph-llm/pyproject.toml b/hugegraph-llm/pyproject.toml index 2b0f29ace..b3894d7e5 100644 --- a/hugegraph-llm/pyproject.toml +++ b/hugegraph-llm/pyproject.toml @@ -17,7 +17,7 @@ [project] name = "hugegraph-llm" -version = "1.5.0" +version = "1.7.0" description = "A tool for the implementation and research related to large language models." authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, @@ -89,4 +89,6 @@ allow-direct-references = true [tool.uv.sources] hugegraph-python-client = { workspace = true } -pycgraph = { git = "https://github.com/ChunelFeng/CGraph.git", subdirectory = "python", rev = "main", marker = "sys_platform == 'linux'" } +# We encountered a bug in PyCGraph's latest release version, so we're using a specific commit from the main branch (without the bug) as the project dependency. +# TODO: Replace this command in the future when a new PyCGraph release version (after 3.1.2) is available. +pycgraph = { git = "https://github.com/ChunelFeng/CGraph.git", subdirectory = "python", rev = "248bfcfeddfa2bc23a1d585a3925c71189dba6cc"} diff --git a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py index 4c192c29c..109da4a99 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/admin_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/admin_api.py @@ -31,7 +31,7 @@ def admin_http_api(router: APIRouter, log_stream): @router.post("/logs", status_code=status.HTTP_200_OK) async def log_stream_api(req: LogStreamRequest): if admin_settings.admin_token != req.admin_token: - raise generate_response( + raise generate_response( # pylint: disable=raising-bad-type RAGResponse( status_code=status.HTTP_403_FORBIDDEN, # pylint: disable=E0702 message="Invalid admin_token", diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py index 356176e4e..ca29cb9ab 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py +++ b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py @@ -31,9 +31,8 @@ from hugegraph_llm.api.models.rag_response import RAGResponse from hugegraph_llm.config import huge_settings from hugegraph_llm.config import llm_settings, prompt +from hugegraph_llm.utils.graph_index_utils import get_vertex_details from hugegraph_llm.utils.log import log -from hugegraph_llm.flows.scheduler import SchedulerSingleton - # pylint: disable=too-many-statements @@ -51,6 +50,13 @@ def rag_http_api( def rag_answer_api(req: RAGRequest): set_graph_config(req) + # Basic parameter validation: empty query => 400 + if not req.query or not str(req.query).strip(): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Query must not be empty.", + ) + result = rag_answer_func( text=req.query, raw_answer=req.raw_answer, @@ -68,7 +74,8 @@ def rag_answer_api(req: RAGRequest): # Keep prompt params in the end custom_related_information=req.custom_priority_info, answer_prompt=req.answer_prompt or prompt.answer_prompt, - keywords_extract_prompt=req.keywords_extract_prompt or prompt.keywords_extract_prompt, + keywords_extract_prompt=req.keywords_extract_prompt + or prompt.keywords_extract_prompt, gremlin_prompt=req.gremlin_prompt or prompt.gremlin_generate_prompt, ) # TODO: we need more info in the response for users to understand the query logic @@ -77,7 +84,8 @@ def rag_answer_api(req: RAGRequest): **{ key: value for key, value in zip( - ["raw_answer", "vector_only", "graph_only", "graph_vector_answer"], result + ["raw_answer", "vector_only", "graph_only", "graph_vector_answer"], + result, ) if getattr(req, key) }, @@ -96,6 +104,13 @@ def graph_rag_recall_api(req: GraphRAGRequest): try: set_graph_config(req) + # Basic parameter validation: empty query => 400 + if not req.query or not str(req.query).strip(): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Query must not be empty.", + ) + result = graph_rag_recall_func( query=req.query, max_graph_items=req.max_graph_items, @@ -111,12 +126,7 @@ def graph_rag_recall_api(req: GraphRAGRequest): ) if req.get_vertex_only: - from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery - - graph_rag = GraphRAGQuery() - graph_rag.init_client(result) - vertex_details = graph_rag.get_vertex_details(result["match_vids"]) - + vertex_details = get_vertex_details(result["match_vids"], result) if vertex_details: result["match_vids"] = vertex_details @@ -136,7 +146,9 @@ def graph_rag_recall_api(req: GraphRAGRequest): except TypeError as e: log.error("TypeError in graph_rag_recall_api: %s", e) - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) + ) from e except Exception as e: log.error("Unexpected error occurred: %s", e) raise HTTPException( @@ -147,7 +159,9 @@ def graph_rag_recall_api(req: GraphRAGRequest): @router.post("/config/graph", status_code=status.HTTP_201_CREATED) def graph_config_api(req: GraphConfigRequest): # Accept status code - res = apply_graph_conf(req.url, req.name, req.user, req.pwd, req.gs, origin_call="http") + res = apply_graph_conf( + req.url, req.name, req.user, req.pwd, req.gs, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) # TODO: restructure the implement of llm to three types, like "/config/chat_llm" @@ -157,10 +171,16 @@ def llm_config_api(req: LLMConfigRequest): if req.llm_type == "openai": res = apply_llm_conf( - req.api_key, req.api_base, req.language_model, req.max_tokens, origin_call="http" + req.api_key, + req.api_base, + req.language_model, + req.max_tokens, + origin_call="http", ) else: - res = apply_llm_conf(req.host, req.port, req.language_model, None, origin_call="http") + res = apply_llm_conf( + req.host, req.port, req.language_model, None, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/embedding", status_code=status.HTTP_201_CREATED) @@ -172,7 +192,9 @@ def embedding_config_api(req: LLMConfigRequest): req.api_key, req.api_base, req.language_model, origin_call="http" ) else: - res = apply_embedding_conf(req.host, req.port, req.language_model, origin_call="http") + res = apply_embedding_conf( + req.host, req.port, req.language_model, origin_call="http" + ) return generate_response(RAGResponse(status_code=res, message="Missing Value")) @router.post("/config/rerank", status_code=status.HTTP_201_CREATED) @@ -184,7 +206,9 @@ def rerank_config_api(req: RerankerConfigRequest): req.api_key, req.reranker_model, req.cohere_base_url, origin_call="http" ) elif req.reranker_type == "siliconflow": - res = apply_reranker_conf(req.api_key, req.reranker_model, None, origin_call="http") + res = apply_reranker_conf( + req.api_key, req.reranker_model, None, origin_call="http" + ) else: res = status.HTTP_501_NOT_IMPLEMENTED return generate_response(RAGResponse(status_code=res, message="Missing Value")) @@ -197,20 +221,20 @@ def text2gremlin_api(req: GremlinGenerateRequest): # Basic parameter validation: empty query => 400 if not req.query or not str(req.query).strip(): raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, detail="Query must not be empty." + status_code=status.HTTP_400_BAD_REQUEST, + detail="Query must not be empty.", ) output_types_str_list = None if req.output_types: output_types_str_list = [ot.value for ot in req.output_types] - response_dict = SchedulerSingleton.get_instance().schedule_flow( - "text2gremlin", - req.query, - req.example_num, - huge_settings.graph_name, - req.gremlin_prompt, - output_types_str_list, + response_dict = gremlin_generate_selective_func( + inp=req.query, + example_num=req.example_num, + schema_input=huge_settings.graph_name, + gremlin_prompt_input=req.gremlin_prompt, + requested_outputs=output_types_str_list, ) return response_dict except HTTPException as e: diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index ca36867d9..9bf04b570 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py @@ -19,12 +19,12 @@ import os from typing import AsyncGenerator, Literal, Optional, Tuple - -import gradio as gr -from hugegraph_llm.flows.scheduler import SchedulerSingleton import pandas as pd +import gradio as gr from gradio.utils import NamedString +from hugegraph_llm.flows import FlowName +from hugegraph_llm.flows.scheduler import SchedulerSingleton from hugegraph_llm.config import resource_path, prompt, llm_settings from hugegraph_llm.utils.decorators import with_task_id from hugegraph_llm.utils.log import log @@ -51,11 +51,7 @@ def rag_answer( ) -> Tuple: """ Generate an answer using the RAG (Retrieval-Augmented Generation) pipeline. - 1. Initialize the RAGPipeline. - 2. Select vector search or graph search based on parameters. - 3. Merge, deduplicate, and rerank the results. - 4. Synthesize the final answer. - 5. Run the pipeline and return the results. + Fetch the Scheduler to deal with the request """ graph_search, gremlin_prompt, vector_search = update_ui_configs( answer_prompt, @@ -75,13 +71,13 @@ def rag_answer( try: # Select workflow by mode to avoid fetching the wrong pipeline from the pool if graph_vector_answer or (graph_only_answer and vector_only_answer): - flow_key = "rag_graph_vector" + flow_key = FlowName.RAG_GRAPH_VECTOR elif vector_only_answer: - flow_key = "rag_vector_only" + flow_key = FlowName.RAG_VECTOR_ONLY elif graph_only_answer: - flow_key = "rag_graph_only" + flow_key = FlowName.RAG_GRAPH_ONLY elif raw_answer: - flow_key = "rag_raw" + flow_key = FlowName.RAG_RAW else: raise RuntimeError("Unsupported flow type") @@ -172,11 +168,7 @@ async def rag_answer_streaming( ) -> AsyncGenerator[Tuple[str, str, str, str], None]: """ Generate an answer using the RAG (Retrieval-Augmented Generation) pipeline. - 1. Initialize the RAGPipeline. - 2. Select vector search or graph search based on parameters. - 3. Merge, deduplicate, and rerank the results. - 4. Synthesize the final answer. - 5. Run the pipeline and return the results. + Fetch the Scheduler to deal with the request """ graph_search, gremlin_prompt, vector_search = update_ui_configs( answer_prompt, @@ -197,13 +189,13 @@ async def rag_answer_streaming( # Select the specific streaming workflow scheduler = SchedulerSingleton.get_instance() if graph_vector_answer or (graph_only_answer and vector_only_answer): - flow_key = "rag_graph_vector" + flow_key = FlowName.RAG_GRAPH_VECTOR elif vector_only_answer: - flow_key = "rag_vector_only" + flow_key = FlowName.RAG_VECTOR_ONLY elif graph_only_answer: - flow_key = "rag_graph_only" + flow_key = FlowName.RAG_GRAPH_ONLY elif raw_answer: - flow_key = "rag_raw" + flow_key = FlowName.RAG_RAW else: raise RuntimeError("Unsupported flow type") @@ -367,7 +359,7 @@ def toggle_slider(enable): ) gr.Markdown( - """## 2. (Batch) Back-testing ) + """## 2. (Batch) Back-testing > 1. Download the template file & fill in the questions you want to test. > 2. Upload the file & click the button to generate answers. (Preview shows the first 40 lines) > 3. The answer options are the same as the above RAG/Q&A frame diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py index 6600d7c41..aa9c2f0c5 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py @@ -25,11 +25,7 @@ import pandas as pd from hugegraph_llm.config import prompt, resource_path, huge_settings -from hugegraph_llm.models.embeddings.init_embedding import Embeddings -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.graph_rag_task import RAGPipeline -from hugegraph_llm.operators.gremlin_generate_task import GremlinGenerator -from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager +from hugegraph_llm.flows import FlowName from hugegraph_llm.utils.embedding_utils import get_index_folder_name from hugegraph_llm.utils.hugegraph_utils import run_gremlin_query from hugegraph_llm.utils.log import log @@ -86,7 +82,9 @@ def store_schema(schema, question, gremlin_prompt): def build_example_vector_index(temp_file) -> dict: - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) index_path = os.path.join(resource_path, folder_name, "gremlin_examples") if not os.path.exists(index_path): os.makedirs(index_path) @@ -98,7 +96,9 @@ def build_example_vector_index(temp_file) -> dict: timestamp = datetime.now().strftime("%Y%m%d%H%M%S") _, file_name = os.path.split(f"{name}_{timestamp}{ext}") log.info("Copying file to: %s", file_name) - target_file = os.path.join(resource_path, folder_name, "gremlin_examples", file_name) + target_file = os.path.join( + resource_path, folder_name, "gremlin_examples", file_name + ) try: import shutil @@ -116,11 +116,10 @@ def build_example_vector_index(temp_file) -> dict: else: log.critical("Unsupported file format. Please input a JSON or CSV file.") return {"error": "Unsupported file format. Please input a JSON or CSV file."} - builder = GremlinGenerator( - llm=LLMs().get_text2gql_llm(), - embedding=Embeddings().get_embedding(), + + return SchedulerSingleton.get_instance().schedule_flow( + FlowName.BUILD_EXAMPLES_INDEX, examples ) - return builder.example_index_build(examples).run() def _process_schema(schema, generator, sm): @@ -182,43 +181,6 @@ def _execute_queries(context, output_types): context["raw_exec_res"] = "" -def gremlin_generate( - inp, example_num, schema, gremlin_prompt, requested_outputs: Optional[List[str]] = None -) -> GremlinResult: - generator = GremlinGenerator( - llm=LLMs().get_text2gql_llm(), embedding=Embeddings().get_embedding() - ) - sm = SchemaManager(graph_name=schema) - - processed_schema, short_schema = _process_schema(schema, generator, sm) - if processed_schema is None and short_schema is None: - return GremlinResult.error("Invalid JSON schema, please check the format carefully.") - - updated_schema = sm.simple_schema(processed_schema) if short_schema else processed_schema - store_schema(str(updated_schema), inp, gremlin_prompt) - - output_types = _configure_output_types(requested_outputs) - - context = ( - generator.example_index_query(example_num) - .gremlin_generate_synthesize(updated_schema, gremlin_prompt) - .run(query=inp) - ) - - _execute_queries(context, output_types) - - match_result = json.dumps( - context.get("match_result", "No Results"), ensure_ascii=False, indent=2 - ) - return GremlinResult.success_result( - match_result=match_result, - template_gremlin=context["result"], - raw_gremlin=context["raw_result"], - template_exec=context["template_exec_res"], - raw_exec=context["raw_exec_res"], - ) - - def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: mini_schema = {} @@ -226,7 +188,11 @@ def simple_schema(schema: Dict[str, Any]) -> Dict[str, Any]: if "vertexlabels" in schema: mini_schema["vertexlabels"] = [] for vertex in schema["vertexlabels"]: - new_vertex = {key: vertex[key] for key in ["id", "name", "properties"] if key in vertex} + new_vertex = { + key: vertex[key] + for key in ["id", "name", "properties"] + if key in vertex + } mini_schema["vertexlabels"].append(new_vertex) # Add necessary edgelabels items (4) @@ -248,7 +214,7 @@ def gremlin_generate_for_ui(inp, example_num, schema, gremlin_prompt): # Execute via scheduler try: res = SchedulerSingleton.get_instance().schedule_flow( - "text2gremlin", + FlowName.TEXT2GREMLIN, inp, int(example_num) if isinstance(example_num, (int, float, str)) else 2, schema, @@ -305,15 +271,21 @@ def create_text2gremlin_block() -> Tuple: with gr.Row(): with gr.Column(scale=1): input_box = gr.Textbox( - value=prompt.default_question, label="Nature Language Query", show_copy_button=True + value=prompt.default_question, + label="Nature Language Query", + show_copy_button=True, ) match = gr.Code( label="Similar Template (TopN)", language="javascript", elem_classes="code-container-show", ) - initialized_out = gr.Textbox(label="Gremlin With Template", show_copy_button=True) - raw_out = gr.Textbox(label="Gremlin Without Template", show_copy_button=True) + initialized_out = gr.Textbox( + label="Gremlin With Template", show_copy_button=True + ) + raw_out = gr.Textbox( + label="Gremlin Without Template", show_copy_button=True + ) tmpl_exec_out = gr.Code( label="Query With Template Output", language="json", @@ -330,7 +302,10 @@ def create_text2gremlin_block() -> Tuple: minimum=0, maximum=10, step=1, value=2, label="Number of refer examples" ) schema_box = gr.Textbox( - value=prompt.text2gql_graph_schema, label="Schema", lines=2, show_copy_button=True + value=prompt.text2gql_graph_schema, + label="Schema", + lines=2, + show_copy_button=True, ) prompt_box = gr.Textbox( value=prompt.gremlin_generate_prompt, @@ -362,24 +337,21 @@ def graph_rag_recall( get_vertex_only: bool = False, ) -> dict: store_schema(prompt.text2gql_graph_schema, query, gremlin_prompt) - rag = RAGPipeline() - rag.extract_keywords().keywords_to_vid( + context = SchedulerSingleton.get_instance().schedule_flow( + FlowName.RAG_GRAPH_ONLY, + query=query, + gremlin_tmpl_num=gremlin_tmpl_num, + rerank_method=rerank_method, + near_neighbor_first=near_neighbor_first, + custom_related_information=custom_related_information, + gremlin_prompt=gremlin_prompt, + max_graph_items=max_graph_items, + topk_return_results=topk_return_results, vector_dis_threshold=vector_dis_threshold, topk_per_keyword=topk_per_keyword, + is_graph_rag_recall=True, + is_vector_only=get_vertex_only, ) - - if not get_vertex_only: - rag.import_schema(huge_settings.graph_name).query_graphdb( - num_gremlin_generate_example=gremlin_tmpl_num, - gremlin_prompt=gremlin_prompt, - max_graph_items=max_graph_items, - ).merge_dedup_rerank( - rerank_method=rerank_method, - near_neighbor_first=near_neighbor_first, - custom_related_information=custom_related_information, - topk_return_results=topk_return_results, - ) - context = rag.run(verbose=True, query=query, graph_search=True) return context @@ -390,45 +362,13 @@ def gremlin_generate_selective( gremlin_prompt_input: str, requested_outputs: Optional[List[str]] = None, ) -> Dict[str, Any]: - """ - Wraps the gremlin_generate function to return a dictionary of outputs - based on the requested_outputs list of strings. - """ - output_keys = [ - "match_result", - "template_gremlin", - "raw_gremlin", - "template_execution_result", - "raw_execution_result", - ] - if not requested_outputs: # None or empty list - requested_outputs = output_keys - - result = gremlin_generate( - inp, example_num, schema_input, gremlin_prompt_input, requested_outputs + response_dict = SchedulerSingleton.get_instance().schedule_flow( + FlowName.TEXT2GREMLIN, + inp, + example_num, + schema_input, + gremlin_prompt_input, + requested_outputs, ) - outputs_dict: Dict[str, Any] = {} - - if not result.success: - # Handle error case - if "match_result" in requested_outputs: - outputs_dict["match_result"] = result.match_result - if result.error_message: - outputs_dict["error_detail"] = result.error_message - return outputs_dict - - # Handle successful case - output_mapping = { - "match_result": result.match_result, - "template_gremlin": result.template_gremlin, - "raw_gremlin": result.raw_gremlin, - "template_execution_result": result.template_exec_result, - "raw_execution_result": result.raw_exec_result, - } - - for key in requested_outputs: - if key in output_mapping: - outputs_dict[key] = output_mapping[key] - - return outputs_dict + return response_dict diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py index 56b5de4b3..84d60df7e 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/vector_graph_block.py @@ -26,6 +26,7 @@ from hugegraph_llm.config import huge_settings from hugegraph_llm.config import prompt from hugegraph_llm.config import resource_path +from hugegraph_llm.flows import FlowName from hugegraph_llm.flows.scheduler import SchedulerSingleton from hugegraph_llm.utils.graph_index_utils import ( get_graph_index_info, @@ -63,12 +64,16 @@ def generate_prompt_for_ui(source_text, scenario, example_name): Handles the UI logic for generating a new prompt using the new workflow architecture. """ if not all([source_text, scenario, example_name]): - gr.Warning("Please provide original text, expected scenario, and select an example!") + gr.Warning( + "Please provide original text, expected scenario, and select an example!" + ) return gr.update() try: # using new architecture scheduler = SchedulerSingleton.get_instance() - result = scheduler.schedule_flow("prompt_generate", source_text, scenario, example_name) + result = scheduler.schedule_flow( + FlowName.PROMPT_GENERATE, source_text, scenario, example_name + ) gr.Info("Prompt generated successfully!") return result except Exception as e: @@ -79,7 +84,9 @@ def generate_prompt_for_ui(source_text, scenario, example_name): def load_example_names(): """Load all candidate examples""" try: - examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "prompt_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return [example.get("name", "Unnamed example") for example in examples] @@ -100,16 +107,22 @@ def load_query_examples(): ), ) if language.upper() == "CN": - examples_path = os.path.join(resource_path, "prompt_examples", "query_examples_CN.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "query_examples_CN.json" + ) else: - examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "query_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) except (FileNotFoundError, json.JSONDecodeError): try: - examples_path = os.path.join(resource_path, "prompt_examples", "query_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "query_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -120,7 +133,9 @@ def load_query_examples(): def load_schema_fewshot_examples(): """Load few-shot examples from a JSON file""" try: - examples_path = os.path.join(resource_path, "prompt_examples", "schema_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "schema_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: examples = json.load(f) return json.dumps(examples, indent=2, ensure_ascii=False) @@ -131,10 +146,14 @@ def load_schema_fewshot_examples(): def update_example_preview(example_name): """Update the display content based on the selected example name.""" try: - examples_path = os.path.join(resource_path, "prompt_examples", "prompt_examples.json") + examples_path = os.path.join( + resource_path, "prompt_examples", "prompt_examples.json" + ) with open(examples_path, "r", encoding="utf-8") as f: all_examples = json.load(f) - selected_example = next((ex for ex in all_examples if ex.get("name") == example_name), None) + selected_example = next( + (ex for ex in all_examples if ex.get("name") == example_name), None + ) if selected_example: return ( @@ -179,7 +198,9 @@ def _create_prompt_helper_block(demo, input_text, info_extract_template): interactive=False, ) - generate_prompt_btn = gr.Button("🚀 Auto-generate Graph Extract Prompt", variant="primary") + generate_prompt_btn = gr.Button( + "🚀 Auto-generate Graph Extract Prompt", variant="primary" + ) # Bind the change event of the dropdown menu few_shot_dropdown.change( fn=update_example_preview, @@ -271,7 +292,9 @@ def create_vector_graph_block(): lines=15, max_lines=29, ) - out = gr.Code(label="Output Info", language="json", elem_classes="code-container-edit") + out = gr.Code( + label="Output Info", language="json", elem_classes="code-container-edit" + ) with gr.Row(): with gr.Accordion("Get RAG Info", open=False): @@ -280,8 +303,12 @@ def create_vector_graph_block(): graph_index_btn0 = gr.Button("Get Graph Index Info", size="sm") with gr.Accordion("Clear RAG Data", open=False): with gr.Column(): - vector_index_btn1 = gr.Button("Clear Chunks Vector Index", size="sm") - graph_index_btn1 = gr.Button("Clear Graph Vid Vector Index", size="sm") + vector_index_btn1 = gr.Button( + "Clear Chunks Vector Index", size="sm" + ) + graph_index_btn1 = gr.Button( + "Clear Graph Vid Vector Index", size="sm" + ) graph_data_btn0 = gr.Button("Clear Graph Data", size="sm") vector_import_bt = gr.Button("Import into Vector", variant="primary") @@ -354,9 +381,9 @@ def create_vector_graph_block(): inputs=[input_text, input_schema, info_extract_template], ) - graph_loading_bt.click(import_graph_data, inputs=[out, input_schema], outputs=[out]).then( - update_vid_embedding - ).then( + graph_loading_bt.click( + import_graph_data, inputs=[out, input_schema], outputs=[out] + ).then(update_vid_embedding).then( store_prompt, inputs=[input_text, input_schema, info_extract_template], ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/__init__.py b/hugegraph-llm/src/hugegraph_llm/flows/__init__.py index 13a83393a..1016680b5 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/__init__.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/__init__.py @@ -14,3 +14,21 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. + +from enum import Enum + + +class FlowName(str, Enum): + RAG_GRAPH_ONLY = "rag_graph_only" + RAG_VECTOR_ONLY = "rag_vector_only" + TEXT2GREMLIN = "text2gremlin" + BUILD_EXAMPLES_INDEX = "build_examples_index" + BUILD_VECTOR_INDEX = "build_vector_index" + GRAPH_EXTRACT = "graph_extract" + IMPORT_GRAPH_DATA = "import_graph_data" + UPDATE_VID_EMBEDDINGS = "update_vid_embeddings" + GET_GRAPH_INDEX_INFO = "get_graph_index_info" + BUILD_SCHEMA = "build_schema" + PROMPT_GENERATE = "prompt_generate" + RAG_RAW = "rag_raw" + RAG_GRAPH_VECTOR = "rag_graph_vector" diff --git a/hugegraph-llm/src/hugegraph_llm/flows/build_example_index.py b/hugegraph-llm/src/hugegraph_llm/flows/build_example_index.py new file mode 100644 index 000000000..d09cc7828 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/build_example_index.py @@ -0,0 +1,62 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from typing import List, Dict, Optional + +from PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.nodes.index_node.build_gremlin_example_index import ( + BuildGremlinExampleIndexNode, +) +from hugegraph_llm.utils.log import log + + +# pylint: disable=arguments-differ,keyword-arg-before-vararg +class BuildExampleIndexFlow(BaseFlow): + def __init__(self): + pass + + def prepare( + self, + prepared_input: WkFlowInput, + examples: Optional[List[Dict[str, str]]], + **kwargs, + ): + prepared_input.examples = examples + + def build_flow(self, examples=None, **kwargs): + pipeline = GPipeline() + prepared_input = WkFlowInput() + self.prepare(prepared_input, examples=examples) + + pipeline.createGParam(prepared_input, "wkflow_input") + pipeline.createGParam(WkFlowState(), "wkflow_state") + + build_node = BuildGremlinExampleIndexNode() + pipeline.registerGElement(build_node, set(), "build_examples_index") + + return pipeline + + def post_deal(self, pipeline=None, **kwargs): + state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + try: + formatted_schema = json.dumps(state_json, ensure_ascii=False, indent=2) + return formatted_schema + except (TypeError, ValueError) as e: + log.error("Failed to format schema: %s", e) + return str(state_json) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/build_schema.py b/hugegraph-llm/src/hugegraph_llm/flows/build_schema.py index 6bbcb8512..1554e53fe 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/build_schema.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/build_schema.py @@ -13,15 +13,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json + +from PyCGraph import GPipeline + from hugegraph_llm.flows.common import BaseFlow from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState from hugegraph_llm.nodes.llm_node.schema_build import SchemaBuildNode from hugegraph_llm.utils.log import log -import json -from PyCGraph import GPipeline - +# pylint: disable=arguments-differ,keyword-arg-before-vararg class BuildSchemaFlow(BaseFlow): def __init__(self): pass @@ -32,15 +34,17 @@ def prepare( texts=None, query_examples=None, few_shot_schema=None, + **kwargs, ): prepared_input.texts = texts # Optional fields packed into wk_input for SchemaBuildNode # Keep raw values; node will parse if strings prepared_input.query_examples = query_examples prepared_input.few_shot_schema = few_shot_schema - return - def build_flow(self, texts=None, query_examples=None, few_shot_schema=None): + def build_flow( + self, texts=None, query_examples=None, few_shot_schema=None, **kwargs + ): pipeline = GPipeline() prepared_input = WkFlowInput() self.prepare( @@ -58,7 +62,7 @@ def build_flow(self, texts=None, query_examples=None, few_shot_schema=None): return pipeline - def post_deal(self, pipeline=None): + def post_deal(self, pipeline=None, **kwargs): state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() if "schema" not in state_json: return "" diff --git a/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py index 9a07b5dba..b57cbfbe3 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/build_vector_index.py @@ -13,28 +13,28 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json + +from PyCGraph import GPipeline + from hugegraph_llm.flows.common import BaseFlow from hugegraph_llm.nodes.document_node.chunk_split import ChunkSplitNode from hugegraph_llm.nodes.index_node.build_vector_index import BuildVectorIndexNode from hugegraph_llm.state.ai_state import WkFlowInput - -import json -from PyCGraph import GPipeline - from hugegraph_llm.state.ai_state import WkFlowState +# pylint: disable=arguments-differ,keyword-arg-before-vararg class BuildVectorIndexFlow(BaseFlow): def __init__(self): pass - def prepare(self, prepared_input: WkFlowInput, texts): + def prepare(self, prepared_input: WkFlowInput, texts, **kwargs): prepared_input.texts = texts prepared_input.language = "zh" prepared_input.split_type = "paragraph" - return - def build_flow(self, texts): + def build_flow(self, texts, **kwargs): pipeline = GPipeline() # prepare for workflow input prepared_input = WkFlowInput() @@ -50,6 +50,6 @@ def build_flow(self, texts): return pipeline - def post_deal(self, pipeline=None): + def post_deal(self, pipeline=None, **kwargs): res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() return json.dumps(res, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/common.py b/hugegraph-llm/src/hugegraph_llm/flows/common.py index e2348466c..d1301119e 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/common.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/common.py @@ -26,25 +26,22 @@ class BaseFlow(ABC): """ @abstractmethod - def prepare(self, prepared_input: WkFlowInput, *args, **kwargs): + def prepare(self, prepared_input: WkFlowInput, **kwargs): """ Pre-processing interface. """ - pass @abstractmethod - def build_flow(self, *args, **kwargs): + def build_flow(self, **kwargs): """ Interface for building the flow. """ - pass @abstractmethod - def post_deal(self, *args, **kwargs): + def post_deal(self, **kwargs): """ Post-processing interface. """ - pass async def post_deal_stream( self, pipeline=None @@ -57,15 +54,11 @@ async def post_deal_stream( if pipeline is None: yield {"error": "No pipeline provided"} return - try: - state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - log.info(f"{flow_name} post processing success") - stream_flow = state_json.get("stream_generator") - if stream_flow is None: - yield {"error": "No stream_generator found in workflow state"} - return - async for chunk in stream_flow: - yield chunk - except Exception as e: - log.error(f"{flow_name} post processing failed: {e}") - yield {"error": f"Post processing failed: {str(e)}"} + state_json = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("%s post processing success", flow_name) + stream_flow = state_json.get("stream_generator") + if stream_flow is None: + yield {"error": "No stream_generator found in workflow state"} + return + async for chunk in stream_flow: + yield chunk diff --git a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py index 7d2735352..86d08bf2d 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py @@ -16,30 +16,32 @@ import json import os +from PyCGraph import GPipeline + from hugegraph_llm.config import huge_settings, llm_settings, resource_path from hugegraph_llm.flows.common import BaseFlow from hugegraph_llm.indices.vector_index import VectorIndex from hugegraph_llm.models.embeddings.init_embedding import model_map from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState from hugegraph_llm.nodes.hugegraph_node.fetch_graph_data import FetchGraphDataNode -from PyCGraph import GPipeline from hugegraph_llm.utils.embedding_utils import ( get_filename_prefix, get_index_folder_name, ) +# pylint: disable=arguments-differ,keyword-arg-before-vararg class GetGraphIndexInfoFlow(BaseFlow): def __init__(self): pass - def prepare(self, prepared_input: WkFlowInput, *args, **kwargs): + def prepare(self, prepared_input: WkFlowInput, **kwargs): return - def build_flow(self, *args, **kwargs): + def build_flow(self, **kwargs): pipeline = GPipeline() prepared_input = WkFlowInput() - self.prepare(prepared_input, *args, **kwargs) + self.prepare(prepared_input, **kwargs) pipeline.createGParam(prepared_input, "wkflow_input") pipeline.createGParam(WkFlowState(), "wkflow_state") fetch_node = FetchGraphDataNode() @@ -48,7 +50,9 @@ def build_flow(self, *args, **kwargs): def post_deal(self, pipeline=None): graph_summary_info = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) filename_prefix = get_filename_prefix( llm_settings.embedding_type, @@ -56,7 +60,7 @@ def post_deal(self, pipeline=None): ) try: vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) - except FileNotFoundError: + except (RuntimeError, OSError): return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) graph_summary_info["vid_index"] = { "embed_dim": vector_index.index.d, diff --git a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py index 55f53b7ad..f3d166786 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/graph_extract.py @@ -23,25 +23,38 @@ from hugegraph_llm.utils.log import log +# pylint: disable=arguments-differ,keyword-arg-before-vararg class GraphExtractFlow(BaseFlow): def __init__(self): pass - def prepare(self, prepared_input: WkFlowInput, schema, texts, example_prompt, extract_type): + def prepare( + self, + prepared_input: WkFlowInput, + schema, + texts, + example_prompt, + extract_type, + language="zh", + **kwargs, + ): # prepare input data prepared_input.texts = texts - prepared_input.language = "zh" + prepared_input.language = language prepared_input.split_type = "document" prepared_input.example_prompt = example_prompt prepared_input.schema = schema prepared_input.extract_type = extract_type - return - def build_flow(self, schema, texts, example_prompt, extract_type): + def build_flow( + self, schema, texts, example_prompt, extract_type, language="zh", **kwargs + ): pipeline = GPipeline() prepared_input = WkFlowInput() # prepare input data - self.prepare(prepared_input, schema, texts, example_prompt, extract_type) + self.prepare( + prepared_input, schema, texts, example_prompt, extract_type, language + ) pipeline.createGParam(prepared_input, "wkflow_input") pipeline.createGParam(WkFlowState(), "wkflow_state") @@ -57,7 +70,7 @@ def build_flow(self, schema, texts, example_prompt, extract_type): return pipeline - def post_deal(self, pipeline=None): + def post_deal(self, pipeline=None, **kwargs): res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() vertices = res.get("vertices", []) edges = res.get("edges", []) 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 0b29b4e64..d0e34ac59 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/import_graph_data.py @@ -24,11 +24,12 @@ from hugegraph_llm.utils.log import log +# pylint: disable=arguments-differ,keyword-arg-before-vararg class ImportGraphDataFlow(BaseFlow): def __init__(self): pass - def prepare(self, prepared_input: WkFlowInput, data, schema): + def prepare(self, prepared_input: WkFlowInput, data, schema, **kwargs): try: data_json = json.loads(data.strip()) if isinstance(data, str) else data except json.JSONDecodeError as e: @@ -43,9 +44,8 @@ def prepare(self, prepared_input: WkFlowInput, data, schema): ) prepared_input.data_json = data_json prepared_input.schema = schema - return - def build_flow(self, data, schema): + def build_flow(self, data, schema, **kwargs): pipeline = GPipeline() prepared_input = WkFlowInput() # prepare input data @@ -61,7 +61,7 @@ def build_flow(self, data, schema): return pipeline - def post_deal(self, pipeline=None): + def post_deal(self, pipeline=None, **kwargs): res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() gr.Info("Import graph data successfully!") return json.dumps(res, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py b/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py index b4a7bf329..16618e13a 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/prompt_generate.py @@ -13,29 +13,30 @@ # See the License for the specific language governing permissions and # limitations under the License. +from PyCGraph import GPipeline + from hugegraph_llm.flows.common import BaseFlow from hugegraph_llm.nodes.llm_node.prompt_generate import PromptGenerateNode from hugegraph_llm.state.ai_state import WkFlowInput - -from PyCGraph import GPipeline - from hugegraph_llm.state.ai_state import WkFlowState +# pylint: disable=arguments-differ,keyword-arg-before-vararg class PromptGenerateFlow(BaseFlow): def __init__(self): pass - def prepare(self, prepared_input: WkFlowInput, source_text, scenario, example_name): + def prepare( + self, prepared_input: WkFlowInput, source_text, scenario, example_name, **kwargs + ): """ Prepare input data for PromptGenerate workflow """ prepared_input.source_text = source_text prepared_input.scenario = scenario prepared_input.example_name = example_name - return - def build_flow(self, source_text, scenario, example_name): + def build_flow(self, source_text, scenario, example_name, **kwargs): """ Build the PromptGenerate workflow """ @@ -53,9 +54,11 @@ def build_flow(self, source_text, scenario, example_name): return pipeline - def post_deal(self, pipeline=None): + def post_deal(self, pipeline=None, **kwargs): """ Process the execution result of PromptGenerate workflow """ res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - return res.get("generated_extract_prompt", "Generation failed. Please check the logs.") + return res.get( + "generated_extract_prompt", "Generation failed. Please check the logs." + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py index 5feb3d471..3029b6259 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py @@ -13,11 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json -from typing import Optional, Literal +from typing import Optional, Literal, cast -from PyCGraph import GPipeline +from PyCGraph import GPipeline, GRegion, GCondition from hugegraph_llm.flows.common import BaseFlow from hugegraph_llm.nodes.llm_node.keyword_extract_node import KeywordExtractNode @@ -31,6 +30,23 @@ from hugegraph_llm.utils.log import log +class GraphRecallCondition(GCondition): + def choose(self): + prepared_input: WkFlowInput = cast( + WkFlowInput, self.getGParamWithNoEmpty("wkflow_input") + ) + return 0 if prepared_input.is_graph_rag_recall else 1 + + +class VectorOnlyCondition(GCondition): + def choose(self): + prepared_input: WkFlowInput = cast( + WkFlowInput, self.getGParamWithNoEmpty("wkflow_input") + ) + return 0 if prepared_input.is_vector_only else 1 + + +# pylint: disable=arguments-differ,keyword-arg-before-vararg class RAGGraphOnlyFlow(BaseFlow): """ Workflow for graph-only answering (graph_only_answer) @@ -40,13 +56,12 @@ def prepare( self, prepared_input: WkFlowInput, query: str, - vector_search: bool = None, - graph_search: bool = None, - raw_answer: bool = None, - vector_only_answer: bool = None, - graph_only_answer: bool = None, - graph_vector_answer: bool = None, - graph_ratio: float = 0.5, + vector_search: bool = False, + graph_search: bool = True, + raw_answer: bool = False, + vector_only_answer: bool = False, + graph_only_answer: bool = True, + graph_vector_answer: bool = False, rerank_method: Literal["bleu", "reranker"] = "bleu", near_neighbor_first: bool = False, custom_related_information: str = "", @@ -54,11 +69,13 @@ def prepare( keywords_extract_prompt: Optional[str] = None, gremlin_tmpl_num: Optional[int] = -1, gremlin_prompt: Optional[str] = None, - max_graph_items: int = None, - topk_return_results: int = None, - vector_dis_threshold: float = None, - topk_per_keyword: int = None, - **_: dict, + max_graph_items: Optional[int] = None, + topk_return_results: Optional[int] = None, + vector_dis_threshold: Optional[float] = None, + topk_per_keyword: Optional[int] = None, + is_graph_rag_recall: bool = False, + is_vector_only: bool = False, + **kwargs, ): prepared_input.query = query prepared_input.vector_search = vector_search @@ -90,13 +107,15 @@ def prepare( ) prepared_input.schema = huge_settings.graph_name + prepared_input.is_graph_rag_recall = is_graph_rag_recall + prepared_input.is_vector_only = is_vector_only prepared_input.data_json = { "query": query, "vector_search": vector_search, "graph_search": graph_search, "max_graph_items": max_graph_items or huge_settings.max_graph_items, + "is_graph_rag_recall": is_graph_rag_recall, } - return def build_flow(self, **kwargs): pipeline = GPipeline() @@ -106,48 +125,49 @@ def build_flow(self, **kwargs): pipeline.createGParam(WkFlowState(), "wkflow_state") # Create nodes and register them with registerGElement - only_keyword_extract_node = KeywordExtractNode() - only_semantic_id_query_node = SemanticIdQueryNode() + only_keyword_extract_node = KeywordExtractNode("only_keyword") + only_semantic_id_query_node = SemanticIdQueryNode( + {only_keyword_extract_node}, "only_semantic" + ) + vector_region: GRegion = GRegion( + [only_keyword_extract_node, only_semantic_id_query_node] + ) + only_schema_node = SchemaNode() - only_graph_query_node = GraphQueryNode() - merge_rerank_node = MergeRerankNode() + schema_node = VectorOnlyCondition([GRegion(), only_schema_node]) + only_graph_query_node = GraphQueryNode("only_graph") + merge_rerank_node = MergeRerankNode({only_graph_query_node}, "merge_rerank") + graph_region: GRegion = GRegion([only_graph_query_node, merge_rerank_node]) + graph_condition_region = VectorOnlyCondition([GRegion(), graph_region]) + answer_synthesize_node = AnswerSynthesizeNode() + answer_node = GraphRecallCondition([GRegion(), answer_synthesize_node]) - pipeline.registerGElement(only_keyword_extract_node, set(), "only_keyword") + pipeline.registerGElement(vector_region, set(), "vector_fetch") + pipeline.registerGElement(schema_node, set(), "schema_condition") pipeline.registerGElement( - only_semantic_id_query_node, {only_keyword_extract_node}, "only_semantic" + graph_condition_region, + {schema_node, vector_region}, + "graph_condition", ) - pipeline.registerGElement(only_schema_node, set(), "only_schema") pipeline.registerGElement( - only_graph_query_node, - {only_schema_node, only_semantic_id_query_node}, - "only_graph", + answer_node, {graph_condition_region}, "answer_condition" ) - pipeline.registerGElement( - merge_rerank_node, {only_graph_query_node}, "merge_one" - ) - pipeline.registerGElement(answer_synthesize_node, {merge_rerank_node}, "graph") log.info("RAGGraphOnlyFlow pipeline built successfully") return pipeline - def post_deal(self, pipeline=None): + def post_deal(self, pipeline=None, **kwargs): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) - try: - res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - log.info("RAGGraphOnlyFlow post processing success") - return { + return {"error": "No pipeline provided"} + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGGraphOnlyFlow post processing success") + return ( + { "raw_answer": res.get("raw_answer", ""), "vector_only_answer": res.get("vector_only_answer", ""), "graph_only_answer": res.get("graph_only_answer", ""), "graph_vector_answer": res.get("graph_vector_answer", ""), } - except Exception as e: - log.error(f"RAGGraphOnlyFlow post processing failed: {e}") - return json.dumps( - {"error": f"Post processing failed: {str(e)}"}, - ensure_ascii=False, - indent=2, - ) + if not res.get("is_graph_rag_recall", False) + else res + ) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py index 2f4a2bfa2..96c4ab858 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json from typing import Optional, Literal @@ -32,6 +31,7 @@ from hugegraph_llm.utils.log import log +# pylint: disable=arguments-differ,keyword-arg-before-vararg class RAGGraphVectorFlow(BaseFlow): """ Workflow for graph + vector hybrid answering (graph_vector_answer) @@ -41,12 +41,12 @@ def prepare( self, prepared_input: WkFlowInput, query: str, - vector_search: bool = None, - graph_search: bool = None, - raw_answer: bool = None, - vector_only_answer: bool = None, - graph_only_answer: bool = None, - graph_vector_answer: bool = None, + vector_search: bool = True, + graph_search: bool = True, + raw_answer: bool = False, + vector_only_answer: bool = False, + graph_only_answer: bool = False, + graph_vector_answer: bool = True, graph_ratio: float = 0.5, rerank_method: Literal["bleu", "reranker"] = "bleu", near_neighbor_first: bool = False, @@ -55,11 +55,11 @@ def prepare( keywords_extract_prompt: Optional[str] = None, gremlin_tmpl_num: Optional[int] = -1, gremlin_prompt: Optional[str] = None, - max_graph_items: int = None, - topk_return_results: int = None, - vector_dis_threshold: float = None, - topk_per_keyword: int = None, - **_: dict, + max_graph_items: Optional[int] = None, + topk_return_results: Optional[int] = None, + vector_dis_threshold: Optional[float] = None, + topk_per_keyword: Optional[int] = None, + **kwargs, ): prepared_input.query = query prepared_input.vector_search = vector_search @@ -98,7 +98,6 @@ def prepare( "graph_search": graph_search, "max_graph_items": max_graph_items or huge_settings.max_graph_items, } - return def build_flow(self, **kwargs): pipeline = GPipeline() @@ -135,24 +134,14 @@ def build_flow(self, **kwargs): log.info("RAGGraphVectorFlow pipeline built successfully") return pipeline - def post_deal(self, pipeline=None): + def post_deal(self, pipeline=None, **kwargs): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) - try: - res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - log.info("RAGGraphVectorFlow post processing success") - return { - "raw_answer": res.get("raw_answer", ""), - "vector_only_answer": res.get("vector_only_answer", ""), - "graph_only_answer": res.get("graph_only_answer", ""), - "graph_vector_answer": res.get("graph_vector_answer", ""), - } - except Exception as e: - log.error(f"RAGGraphVectorFlow post processing failed: {e}") - return json.dumps( - {"error": f"Post processing failed: {str(e)}"}, - ensure_ascii=False, - indent=2, - ) + return {"error": "No pipeline provided"} + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGGraphVectorFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py index f62e574bb..ede8f98e1 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_raw.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json from typing import Optional @@ -26,6 +25,7 @@ from hugegraph_llm.utils.log import log +# pylint: disable=arguments-differ,keyword-arg-before-vararg class RAGRawFlow(BaseFlow): """ Workflow for basic LLM answering only (raw_answer) @@ -35,16 +35,16 @@ def prepare( self, prepared_input: WkFlowInput, query: str, - vector_search: bool = None, - graph_search: bool = None, - raw_answer: bool = None, - vector_only_answer: bool = None, - graph_only_answer: bool = None, - graph_vector_answer: bool = None, + vector_search: bool = False, + graph_search: bool = False, + raw_answer: bool = True, + vector_only_answer: bool = False, + graph_only_answer: bool = False, + graph_vector_answer: bool = False, custom_related_information: str = "", answer_prompt: Optional[str] = None, - max_graph_items: int = None, - **_: dict, + max_graph_items: Optional[int] = None, + **kwargs, ): prepared_input.query = query prepared_input.raw_answer = raw_answer @@ -61,7 +61,6 @@ def prepare( "graph_search": graph_search, "max_graph_items": max_graph_items or huge_settings.max_graph_items, } - return def build_flow(self, **kwargs): pipeline = GPipeline() @@ -76,24 +75,14 @@ def build_flow(self, **kwargs): log.info("RAGRawFlow pipeline built successfully") return pipeline - def post_deal(self, pipeline=None): + def post_deal(self, pipeline=None, **kwargs): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) - try: - res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - log.info("RAGRawFlow post processing success") - return { - "raw_answer": res.get("raw_answer", ""), - "vector_only_answer": res.get("vector_only_answer", ""), - "graph_only_answer": res.get("graph_only_answer", ""), - "graph_vector_answer": res.get("graph_vector_answer", ""), - } - except Exception as e: - log.error(f"RAGRawFlow post processing failed: {e}") - return json.dumps( - {"error": f"Post processing failed: {str(e)}"}, - ensure_ascii=False, - indent=2, - ) + return {"error": "No pipeline provided"} + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGRawFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py index c727eacce..150e3162a 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_vector_only.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json from typing import Optional, Literal @@ -28,6 +27,7 @@ from hugegraph_llm.utils.log import log +# pylint: disable=arguments-differ,keyword-arg-before-vararg class RAGVectorOnlyFlow(BaseFlow): """ Workflow for vector-only answering (vector_only_answer) @@ -37,20 +37,20 @@ def prepare( self, prepared_input: WkFlowInput, query: str, - vector_search: bool = None, - graph_search: bool = None, - raw_answer: bool = None, - vector_only_answer: bool = None, - graph_only_answer: bool = None, - graph_vector_answer: bool = None, + vector_search: bool = True, + graph_search: bool = False, + raw_answer: bool = False, + vector_only_answer: bool = True, + graph_only_answer: bool = False, + graph_vector_answer: bool = False, rerank_method: Literal["bleu", "reranker"] = "bleu", near_neighbor_first: bool = False, custom_related_information: str = "", answer_prompt: Optional[str] = None, - max_graph_items: int = None, - topk_return_results: int = None, - vector_dis_threshold: float = None, - **_: dict, + max_graph_items: Optional[int] = None, + topk_return_results: Optional[int] = None, + vector_dis_threshold: Optional[float] = None, + **kwargs, ): prepared_input.query = query prepared_input.vector_search = vector_search @@ -77,7 +77,6 @@ def prepare( "graph_search": graph_search, "max_graph_items": max_graph_items or huge_settings.max_graph_items, } - return def build_flow(self, **kwargs): pipeline = GPipeline() @@ -100,24 +99,14 @@ def build_flow(self, **kwargs): log.info("RAGVectorOnlyFlow pipeline built successfully") return pipeline - def post_deal(self, pipeline=None): + def post_deal(self, pipeline=None, **kwargs): if pipeline is None: - return json.dumps( - {"error": "No pipeline provided"}, ensure_ascii=False, indent=2 - ) - try: - res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - log.info("RAGVectorOnlyFlow post processing success") - return { - "raw_answer": res.get("raw_answer", ""), - "vector_only_answer": res.get("vector_only_answer", ""), - "graph_only_answer": res.get("graph_only_answer", ""), - "graph_vector_answer": res.get("graph_vector_answer", ""), - } - except Exception as e: - log.error(f"RAGVectorOnlyFlow post processing failed: {e}") - return json.dumps( - {"error": f"Post processing failed: {str(e)}"}, - ensure_ascii=False, - indent=2, - ) + return {"error": "No pipeline provided"} + res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() + log.info("RAGVectorOnlyFlow post processing success") + return { + "raw_answer": res.get("raw_answer", ""), + "vector_only_answer": res.get("vector_only_answer", ""), + "graph_only_answer": res.get("graph_only_answer", ""), + "graph_vector_answer": res.get("graph_vector_answer", ""), + } diff --git a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py index 5afa1bf8e..bdf59d84e 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/scheduler.py @@ -16,11 +16,13 @@ import threading from typing import Dict, Any from PyCGraph import GPipeline, GPipelineManager +from hugegraph_llm.flows import FlowName from hugegraph_llm.flows.build_vector_index import BuildVectorIndexFlow from hugegraph_llm.flows.common import BaseFlow +from hugegraph_llm.flows.build_example_index import BuildExampleIndexFlow from hugegraph_llm.flows.graph_extract import GraphExtractFlow from hugegraph_llm.flows.import_graph_data import ImportGraphDataFlow -from hugegraph_llm.flows.update_vid_embeddings import UpdateVidEmbeddingsFlows +from hugegraph_llm.flows.update_vid_embeddings import UpdateVidEmbeddingsFlow from hugegraph_llm.flows.get_graph_index_info import GetGraphIndexInfoFlow from hugegraph_llm.flows.build_schema import BuildSchemaFlow from hugegraph_llm.flows.prompt_generate import PromptGenerateFlow @@ -34,72 +36,76 @@ class Scheduler: - pipeline_pool: Dict[str, Any] = None + pipeline_pool: Dict[str, Any] max_pipeline: int def __init__(self, max_pipeline: int = 10): self.pipeline_pool = {} # pipeline_pool act as a manager of GPipelineManager which used for pipeline management - self.pipeline_pool["build_vector_index"] = { + self.pipeline_pool[FlowName.BUILD_VECTOR_INDEX] = { "manager": GPipelineManager(), "flow": BuildVectorIndexFlow(), } - self.pipeline_pool["graph_extract"] = { + self.pipeline_pool[FlowName.GRAPH_EXTRACT] = { "manager": GPipelineManager(), "flow": GraphExtractFlow(), } - self.pipeline_pool["import_graph_data"] = { + self.pipeline_pool[FlowName.IMPORT_GRAPH_DATA] = { "manager": GPipelineManager(), "flow": ImportGraphDataFlow(), } - self.pipeline_pool["update_vid_embeddings"] = { + self.pipeline_pool[FlowName.UPDATE_VID_EMBEDDINGS] = { "manager": GPipelineManager(), - "flow": UpdateVidEmbeddingsFlows(), + "flow": UpdateVidEmbeddingsFlow(), } - self.pipeline_pool["get_graph_index_info"] = { + self.pipeline_pool[FlowName.GET_GRAPH_INDEX_INFO] = { "manager": GPipelineManager(), "flow": GetGraphIndexInfoFlow(), } - self.pipeline_pool["build_schema"] = { + self.pipeline_pool[FlowName.BUILD_SCHEMA] = { "manager": GPipelineManager(), "flow": BuildSchemaFlow(), } - self.pipeline_pool["prompt_generate"] = { + self.pipeline_pool[FlowName.PROMPT_GENERATE] = { "manager": GPipelineManager(), "flow": PromptGenerateFlow(), } - self.pipeline_pool["text2gremlin"] = { + self.pipeline_pool[FlowName.TEXT2GREMLIN] = { "manager": GPipelineManager(), "flow": Text2GremlinFlow(), } # New split rag pipelines - self.pipeline_pool["rag_raw"] = { + self.pipeline_pool[FlowName.RAG_RAW] = { "manager": GPipelineManager(), "flow": RAGRawFlow(), } - self.pipeline_pool["rag_vector_only"] = { + self.pipeline_pool[FlowName.RAG_VECTOR_ONLY] = { "manager": GPipelineManager(), "flow": RAGVectorOnlyFlow(), } - self.pipeline_pool["rag_graph_only"] = { + self.pipeline_pool[FlowName.RAG_GRAPH_ONLY] = { "manager": GPipelineManager(), "flow": RAGGraphOnlyFlow(), } - self.pipeline_pool["rag_graph_vector"] = { + self.pipeline_pool[FlowName.RAG_GRAPH_VECTOR] = { "manager": GPipelineManager(), "flow": RAGGraphVectorFlow(), } + self.pipeline_pool[FlowName.BUILD_EXAMPLES_INDEX] = { + "manager": GPipelineManager(), + "flow": BuildExampleIndexFlow(), + } self.max_pipeline = max_pipeline # TODO: Implement Agentic Workflow def agentic_flow(self): pass - def schedule_flow(self, flow: str, *args, **kwargs): - if flow not in self.pipeline_pool: - raise ValueError(f"Unsupported workflow {flow}") - manager: GPipelineManager = self.pipeline_pool[flow]["manager"] - flow: BaseFlow = self.pipeline_pool[flow]["flow"] + def schedule_flow(self, flow_name: str, *args, **kwargs): + if flow_name not in self.pipeline_pool: + raise ValueError(f"Unsupported workflow {flow_name}") + manager: GPipelineManager = self.pipeline_pool[flow_name]["manager"] + flow: BaseFlow = self.pipeline_pool[flow_name]["flow"] pipeline: GPipeline = manager.fetch() if pipeline is None: # call coresponding flow_func to create new workflow @@ -111,13 +117,14 @@ def schedule_flow(self, flow: str, *args, **kwargs): raise RuntimeError(error_msg) status = pipeline.run() if status.isErr(): + manager.add(pipeline) error_msg = f"Error in flow execution: {status.getInfo()}" log.error(error_msg) raise RuntimeError(error_msg) res = flow.post_deal(pipeline) manager.add(pipeline) return res - else: + try: # fetch pipeline & prepare input for flow prepared_input = pipeline.getGParamWithNoEmpty("wkflow_input") flow.prepare(prepared_input, *args, **kwargs) @@ -127,49 +134,46 @@ def schedule_flow(self, flow: str, *args, **kwargs): log.error(error_msg) raise RuntimeError(error_msg) res = flow.post_deal(pipeline) + finally: manager.release(pipeline) - return res + return res - async def schedule_stream_flow(self, flow: str, *args, **kwargs): - if flow not in self.pipeline_pool: - raise ValueError(f"Unsupported workflow {flow}") - manager: GPipelineManager = self.pipeline_pool[flow]["manager"] - flow: BaseFlow = self.pipeline_pool[flow]["flow"] + async def schedule_stream_flow(self, flow_name: str, *args, **kwargs): + if flow_name not in self.pipeline_pool: + raise ValueError(f"Unsupported workflow {flow_name}") + manager: GPipelineManager = self.pipeline_pool[flow_name]["manager"] + flow: BaseFlow = self.pipeline_pool[flow_name]["flow"] pipeline: GPipeline = manager.fetch() if pipeline is None: # call coresponding flow_func to create new workflow pipeline = flow.build_flow(*args, **kwargs) - try: - pipeline.getGParamWithNoEmpty("wkflow_input").stream = True - status = pipeline.init() - if status.isErr(): - error_msg = f"Error in flow init: {status.getInfo()}" - log.error(error_msg) - raise RuntimeError(error_msg) - status = pipeline.run() - if status.isErr(): - error_msg = f"Error in flow execution: {status.getInfo()}" - log.error(error_msg) - raise RuntimeError(error_msg) - async for res in flow.post_deal_stream(pipeline): - yield res - finally: + pipeline.getGParamWithNoEmpty("wkflow_input").stream = True + status = pipeline.init() + if status.isErr(): + error_msg = f"Error in flow init: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + status = pipeline.run() + if status.isErr(): manager.add(pipeline) - else: - try: - # fetch pipeline & prepare input for flow - prepared_input: WkFlowInput = pipeline.getGParamWithNoEmpty( - "wkflow_input" - ) - prepared_input.stream = True - flow.prepare(prepared_input, *args, **kwargs) - status = pipeline.run() - if status.isErr(): - raise RuntimeError(f"Error in flow execution {status.getInfo()}") - async for res in flow.post_deal_stream(pipeline): - yield res - finally: - manager.release(pipeline) + error_msg = f"Error in flow execution: {status.getInfo()}" + log.error(error_msg) + raise RuntimeError(error_msg) + async for res in flow.post_deal_stream(pipeline): + yield res + manager.add(pipeline) + try: + # fetch pipeline & prepare input for flow + prepared_input: WkFlowInput = pipeline.getGParamWithNoEmpty("wkflow_input") + prepared_input.stream = True + flow.prepare(prepared_input, *args, **kwargs) + status = pipeline.run() + if status.isErr(): + raise RuntimeError(f"Error in flow execution {status.getInfo()}") + async for res in flow.post_deal_stream(pipeline): + yield res + finally: + manager.release(pipeline) class SchedulerSingleton: diff --git a/hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py index e9ba4276c..1ae5662cb 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/text2gremlin.py @@ -13,18 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any, Dict, List, Optional + from PyCGraph import GPipeline from hugegraph_llm.flows.common import BaseFlow from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode -from hugegraph_llm.nodes.index_node.gremlin_example_index_query import GremlinExampleIndexQueryNode +from hugegraph_llm.nodes.index_node.gremlin_example_index_query import ( + GremlinExampleIndexQueryNode, +) from hugegraph_llm.nodes.llm_node.text2gremlin import Text2GremlinNode from hugegraph_llm.nodes.hugegraph_node.gremlin_execute import GremlinExecuteNode -from typing import Any, Dict, List, Optional - +# pylint: disable=arguments-differ,keyword-arg-before-vararg class Text2GremlinFlow(BaseFlow): def __init__(self): pass @@ -37,6 +40,7 @@ def prepare( schema_input: str, gremlin_prompt_input: Optional[str], requested_outputs: Optional[List[str]], + **kwargs, ): # sanitize example_num to [0,10], fallback to 2 if invalid if not isinstance(example_num, int): @@ -63,7 +67,6 @@ def prepare( prepared_input.schema = schema_input prepared_input.gremlin_prompt = gremlin_prompt_input prepared_input.requested_outputs = req - return def build_flow( self, @@ -72,6 +75,7 @@ def build_flow( schema_input: str, gremlin_prompt_input: Optional[str] = None, requested_outputs: Optional[List[str]] = None, + **kwargs, ): pipeline = GPipeline() @@ -100,7 +104,7 @@ def build_flow( return pipeline - def post_deal(self, pipeline=None) -> Dict[str, Any]: + def post_deal(self, pipeline=None, **kwargs) -> Dict[str, Any]: state = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() # 始终返回 5 个标准键，避免前端因过滤异常看不到字段 return { 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 b3f0d9923..216f35618 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/update_vid_embeddings.py @@ -13,18 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus, GPipeline -from hugegraph_llm.flows.common import BaseFlow, WkFlowInput +from PyCGraph import GPipeline + +from hugegraph_llm.flows.common import BaseFlow from hugegraph_llm.nodes.hugegraph_node.fetch_graph_data import FetchGraphDataNode from hugegraph_llm.nodes.index_node.build_semantic_index import BuildSemanticIndexNode -from hugegraph_llm.state.ai_state import WkFlowState +from hugegraph_llm.state.ai_state import WkFlowState, WkFlowInput -class UpdateVidEmbeddingsFlows(BaseFlow): - def prepare(self, prepared_input: WkFlowInput): - return CStatus() +# pylint: disable=arguments-differ,keyword-arg-before-vararg +class UpdateVidEmbeddingsFlow(BaseFlow): + def prepare(self, prepared_input: WkFlowInput, **kwargs): + pass - def build_flow(self): + def build_flow(self, **kwargs): pipeline = GPipeline() prepared_input = WkFlowInput() # prepare input data @@ -40,7 +42,7 @@ def build_flow(self): return pipeline - def post_deal(self, pipeline): + def post_deal(self, pipeline, **kwargs): res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() removed_num = res.get("removed_vid_vector_num", 0) added_num = res.get("added_vid_vector_num", 0) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/utils.py b/hugegraph-llm/src/hugegraph_llm/flows/utils.py deleted file mode 100644 index b4ba05c84..000000000 --- a/hugegraph-llm/src/hugegraph_llm/flows/utils.py +++ /dev/null @@ -1,34 +0,0 @@ -# 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 hugegraph_llm.state.ai_state import WkFlowInput -from hugegraph_llm.utils.log import log - - -def prepare_schema(prepared_input: WkFlowInput, schema): - schema = schema.strip() - if schema.startswith("{"): - try: - schema = json.loads(schema) - prepared_input.schema = schema - except json.JSONDecodeError as exc: - log.error("Invalid JSON format in schema. Please check it again.") - raise ValueError("Invalid JSON format in schema.") from exc - else: - log.info("Get schema '%s' from graphdb.", schema) - prepared_input.graph_name = schema - return diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index 3ad50b3ec..de04dff87 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -29,27 +29,27 @@ } -def get_embedding(llm_settings: LLMConfig): - if llm_settings.embedding_type == "openai": +def get_embedding(llm_configs: LLMConfig): + if llm_configs.embedding_type == "openai": return OpenAIEmbedding( - model_name=llm_settings.openai_embedding_model, - api_key=llm_settings.openai_embedding_api_key, - api_base=llm_settings.openai_embedding_api_base, + model_name=llm_configs.openai_embedding_model, + api_key=llm_configs.openai_embedding_api_key, + api_base=llm_configs.openai_embedding_api_base, ) - if llm_settings.embedding_type == "ollama/local": + if llm_configs.embedding_type == "ollama/local": return OllamaEmbedding( - model_name=llm_settings.ollama_embedding_model, - host=llm_settings.ollama_embedding_host, - port=llm_settings.ollama_embedding_port, + model_name=llm_configs.ollama_embedding_model, + host=llm_configs.ollama_embedding_host, + port=llm_configs.ollama_embedding_port, ) - if llm_settings.embedding_type == "litellm": + if llm_configs.embedding_type == "litellm": return LiteLLMEmbedding( - model_name=llm_settings.litellm_embedding_model, - api_key=llm_settings.litellm_embedding_api_key, - api_base=llm_settings.litellm_embedding_api_base, + model_name=llm_configs.litellm_embedding_model, + api_key=llm_configs.litellm_embedding_api_key, + api_base=llm_configs.litellm_embedding_api_base, ) - raise Exception("embedding type is not supported !") + raise ValueError("embedding type is not supported !") class Embeddings: 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 9121fca09..a13641db0 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/llms/init_llm.py @@ -22,74 +22,74 @@ from hugegraph_llm.config import llm_settings -def get_chat_llm(llm_settings: LLMConfig): - if llm_settings.chat_llm_type == "openai": +def get_chat_llm(llm_configs: LLMConfig): + if llm_configs.chat_llm_type == "openai": return OpenAIClient( - api_key=llm_settings.openai_chat_api_key, - api_base=llm_settings.openai_chat_api_base, - model_name=llm_settings.openai_chat_language_model, - max_tokens=llm_settings.openai_chat_tokens, + 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, ) - if llm_settings.chat_llm_type == "ollama/local": + if llm_configs.chat_llm_type == "ollama/local": return OllamaClient( - model=llm_settings.ollama_chat_language_model, - host=llm_settings.ollama_chat_host, - port=llm_settings.ollama_chat_port, + model=llm_configs.ollama_chat_language_model, + host=llm_configs.ollama_chat_host, + port=llm_configs.ollama_chat_port, ) - if llm_settings.chat_llm_type == "litellm": + if llm_configs.chat_llm_type == "litellm": return LiteLLMClient( - api_key=llm_settings.litellm_chat_api_key, - api_base=llm_settings.litellm_chat_api_base, - model_name=llm_settings.litellm_chat_language_model, - max_tokens=llm_settings.litellm_chat_tokens, + 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, ) raise Exception("chat llm type is not supported !") -def get_extract_llm(llm_settings: LLMConfig): - if llm_settings.extract_llm_type == "openai": +def get_extract_llm(llm_configs: LLMConfig): + if llm_configs.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, + 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, ) - if llm_settings.extract_llm_type == "ollama/local": + if llm_configs.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, + model=llm_configs.ollama_extract_language_model, + host=llm_configs.ollama_extract_host, + port=llm_configs.ollama_extract_port, ) - if llm_settings.extract_llm_type == "litellm": + if llm_configs.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, + 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, ) raise Exception("extract llm type is not supported !") -def get_text2gql_llm(llm_settings: LLMConfig): - if llm_settings.text2gql_llm_type == "openai": +def get_text2gql_llm(llm_configs: LLMConfig): + if llm_configs.text2gql_llm_type == "openai": return OpenAIClient( - api_key=llm_settings.openai_text2gql_api_key, - api_base=llm_settings.openai_text2gql_api_base, - model_name=llm_settings.openai_text2gql_language_model, - max_tokens=llm_settings.openai_text2gql_tokens, + api_key=llm_configs.openai_text2gql_api_key, + api_base=llm_configs.openai_text2gql_api_base, + model_name=llm_configs.openai_text2gql_language_model, + max_tokens=llm_configs.openai_text2gql_tokens, ) - if llm_settings.text2gql_llm_type == "ollama/local": + if llm_configs.text2gql_llm_type == "ollama/local": return OllamaClient( - model=llm_settings.ollama_text2gql_language_model, - host=llm_settings.ollama_text2gql_host, - port=llm_settings.ollama_text2gql_port, + model=llm_configs.ollama_text2gql_language_model, + host=llm_configs.ollama_text2gql_host, + port=llm_configs.ollama_text2gql_port, ) - if llm_settings.text2gql_llm_type == "litellm": + if llm_configs.text2gql_llm_type == "litellm": return LiteLLMClient( - api_key=llm_settings.litellm_text2gql_api_key, - api_base=llm_settings.litellm_text2gql_api_base, - model_name=llm_settings.litellm_text2gql_language_model, - max_tokens=llm_settings.litellm_text2gql_tokens, + api_key=llm_configs.litellm_text2gql_api_key, + api_base=llm_configs.litellm_text2gql_api_base, + model_name=llm_configs.litellm_text2gql_language_model, + max_tokens=llm_configs.litellm_text2gql_tokens, ) raise Exception("text2gql llm type is not supported !") @@ -173,4 +173,8 @@ def get_text2gql_llm(self): if __name__ == "__main__": client = LLMs().get_chat_llm() print(client.generate(prompt="What is the capital of China?")) - print(client.generate(messages=[{"role": "user", "content": "What is the capital of China?"}])) + print( + client.generate( + messages=[{"role": "user", "content": "What is the capital of China?"}] + ) + ) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py index f90167305..d7e53d4b8 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/base_node.py @@ -13,14 +13,26 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Dict, Optional from PyCGraph import GNode, CStatus from hugegraph_llm.nodes.util import init_context from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState +from hugegraph_llm.utils.log import log class BaseNode(GNode): - context: WkFlowState = None - wk_input: WkFlowInput = None + """ + Base class for workflow nodes, providing context management and operation scheduling. + + All custom nodes should inherit from this class and implement the operator_schedule method. + + Attributes: + context: Shared workflow state + wk_input: Workflow input parameters + """ + + context: Optional[WkFlowState] = None + wk_input: Optional[WkFlowInput] = None def init(self): return init_context(self) @@ -30,6 +42,8 @@ def node_init(self): Node initialization method, can be overridden by subclasses. Returns a CStatus object indicating whether initialization succeeded. """ + if self.wk_input is None or self.context is None: + return CStatus(-1, "wk_input or context not initialized") if self.wk_input.data_json is not None: self.context.assign_from_json(self.wk_input.data_json) self.wk_input.data_json = None @@ -43,6 +57,8 @@ def run(self): sts = self.node_init() if sts.isErr(): return sts + if self.context is None: + return CStatus(-1, "Context not initialized") self.context.lock() try: data_json = self.context.to_json() @@ -51,24 +67,35 @@ def run(self): try: res = self.operator_schedule(data_json) - except Exception as exc: + except (ValueError, TypeError, KeyError, NotImplementedError) as exc: import traceback node_info = f"Node type: {type(self).__name__}, Node object: {self}" err_msg = f"Node failed: {exc}\n{node_info}\n{traceback.format_exc()}" return CStatus(-1, err_msg) + # For unexpected exceptions, re-raise to let them propagate or be caught elsewhere self.context.lock() try: - if isinstance(res, dict): + if res is not None and isinstance(res, dict): self.context.assign_from_json(res) + elif res is not None: + log.warning("operator_schedule returned non-dict type: %s", type(res)) finally: self.context.unlock() return CStatus() - def operator_schedule(self, data_json): + def operator_schedule(self, data_json) -> Optional[Dict]: """ - Interface for scheduling the operator, can be overridden by subclasses. - Returns a CStatus object indicating whether scheduling succeeded. + Operation scheduling method that must be implemented by subclasses. + + Args: + data_json: Context serialized as JSON data + + Returns: + Dictionary of processing results, or None to indicate no update + + Raises: + NotImplementedError: If the subclass has not implemented this method """ - pass + raise NotImplementedError("Subclasses must implement operator_schedule") diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py index 78f53e231..c718086aa 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/common_node/merge_rerank_node.py @@ -52,8 +52,8 @@ def node_init(self): topk_return_results=topk_return_results, ) return super().node_init() - except Exception as e: - log.error(f"Failed to initialize MergeRerankNode: {e}") + except ValueError as e: + log.error("Failed to initialize MergeRerankNode: %s", e) from PyCGraph import CStatus return CStatus(-1, f"MergeRerankNode initialization failed: {e}") @@ -72,12 +72,14 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: merged_count = len(result.get("merged_result", [])) log.info( - f"Merge and rerank completed: {vector_count} vector results, " - f"{graph_count} graph results, {merged_count} merged results" + "Merge and rerank completed: %d vector results, %d graph results, %d merged results", + vector_count, + graph_count, + merged_count, ) return result - except Exception as e: - log.error(f"Merge and rerank failed: {e}") + except ValueError as e: + log.error("Merge and rerank failed: %s", e) return data_json 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 f71bd7bd5..883cc909d 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 @@ -13,8 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from hugegraph_llm.nodes.base_node import BaseNode from PyCGraph import CStatus +from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.document_op.chunk_split import ChunkSplit from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState 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 99b428e5e..6e9dd01ad 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 @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Optional + from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.hugegraph_op.fetch_graph_data import FetchGraphData from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -21,11 +23,12 @@ class FetchGraphDataNode(BaseNode): fetch_graph_data_op: FetchGraphData - context: WkFlowState = None - wk_input: WkFlowInput = None + context: Optional[WkFlowState] = None + wk_input: Optional[WkFlowInput] = None def node_init(self): - self.fetch_graph_data_op = FetchGraphData(get_hg_client()) + client = get_hg_client() + self.fetch_graph_data_op = FetchGraphData(client) return super().node_init() def operator_schedule(self, data_json): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py index ae65ccb33..c9d62a9d5 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/graph_query_node.py @@ -13,12 +13,62 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus -from typing import Dict, Any +import json +from typing import Dict, Any, Tuple, List, Set, Optional + from hugegraph_llm.nodes.base_node import BaseNode -from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.operators.operator_list import OperatorList from hugegraph_llm.utils.log import log +from pyhugegraph.client import PyHugeClient + +# TODO: remove 'as('subj)' step +VERTEX_QUERY_TPL = "g.V({keywords}).limit(8).as('subj').toList()" + +# TODO: we could use a simpler query (like kneighbor-api to get the edges) +# TODO: test with profile()/explain() to speed up the query +VID_QUERY_NEIGHBOR_TPL = """\ +g.V({keywords}) +.repeat( + bothE({edge_labels}).limit({edge_limit}).otherV().dedup() +).times({max_deep}).emit() +.simplePath() +.path() +.by(project('label', 'id', 'props') + .by(label()) + .by(id()) + .by(valueMap().by(unfold())) +) +.by(project('label', 'inV', 'outV', 'props') + .by(label()) + .by(inV().id()) + .by(outV().id()) + .by(valueMap().by(unfold())) +) +.limit({max_items}) +.toList() +""" + +PROPERTY_QUERY_NEIGHBOR_TPL = """\ +g.V().has('{prop}', within({keywords})) +.repeat( + bothE({edge_labels}).limit({edge_limit}).otherV().dedup() +).times({max_deep}).emit() +.simplePath() +.path() +.by(project('label', 'props') + .by(label()) + .by(valueMap().by(unfold())) +) +.by(project('label', 'inV', 'outV', 'props') + .by(label()) + .by(inV().values('{prop}')) + .by(outV().values('{prop}')) + .by(valueMap().by(unfold())) +) +.limit({max_items}) +.toList() +""" class GraphQueryNode(BaseNode): @@ -26,45 +76,395 @@ class GraphQueryNode(BaseNode): Graph query node, responsible for retrieving relevant information from the graph database. """ - graph_rag_query: GraphRAGQuery + _client: Optional[PyHugeClient] = None + _max_deep: Optional[int] = None + _max_items: Optional[int] = None + _prop_to_match: Optional[str] = None + _num_gremlin_generate_example: int = -1 + gremlin_prompt: str = "" + _limit_property: bool = False + _max_v_prop_len: int = 2048 + _max_e_prop_len: int = 256 + _schema: str = "" + operator_list: Optional[OperatorList] = None def node_init(self): """ Initialize the graph query operator. """ + self._client: PyHugeClient = 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, + ) + self._max_deep = self.wk_input.max_deep or 2 + self._max_items = self.wk_input.max_graph_items or huge_settings.max_graph_items + self._prop_to_match = self.wk_input.prop_to_match + self._num_gremlin_generate_example = ( + self.wk_input.gremlin_tmpl_num + if self.wk_input.gremlin_tmpl_num is not None + else -1 + ) + self.gremlin_prompt = ( + self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt + ) + self._limit_property = huge_settings.limit_property.lower() == "true" + self._max_v_prop_len = self.wk_input.max_v_prop_len or 2048 + self._max_e_prop_len = self.wk_input.max_e_prop_len or 256 + self._schema = "" + self.operator_list = OperatorList(None, None) + + return super().node_init() + + # TODO: move this method to a util file for reuse (remove self param) + def init_client(self, context): + """Initialize the HugeGraph client from context or default settings.""" + # pylint: disable=R0915 (too-many-statements) + if self._client is None: + if isinstance(context.get("graph_client"), PyHugeClient): + self._client = context["graph_client"] + else: + url = context.get("url") or "http://localhost:8080" + graph = context.get("graph") or "hugegraph" + user = context.get("user") or "admin" + pwd = context.get("pwd") or "admin" + gs = context.get("graphspace") or None + self._client = PyHugeClient(url, graph, user, pwd, gs) + assert self._client is not None, "No valid graph to search." + + def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: + query = context["query"] + vertices = context.get("match_vids") + query_embedding = context.get("query_embedding") + + self.operator_list.clear() + self.operator_list.example_index_query( + num_examples=self._num_gremlin_generate_example + ) + gremlin_response = self.operator_list.gremlin_generate_synthesize( + context["simple_schema"], + vertices=vertices, + gremlin_prompt=self.gremlin_prompt, + ).run(query=query, query_embedding=query_embedding) + if self._num_gremlin_generate_example > 0: + gremlin = gremlin_response["result"] + else: + gremlin = gremlin_response["raw_result"] + log.info("Generated gremlin: %s", gremlin) + context["gremlin"] = gremlin try: - graph_name = huge_settings.graph_name - if not graph_name: - return CStatus(-1, "graph_name is required in wk_input") + result = self._client.gremlin().exec(gremlin=gremlin)["data"] + if result == [None]: + result = [] + context["graph_result"] = [ + json.dumps(item, ensure_ascii=False) for item in result + ] + if context["graph_result"]: + context["graph_result_flag"] = 1 + context["graph_context_head"] = ( + f"The following are graph query result " + f"from gremlin query `{gremlin}`.\n" + ) + except Exception as e: # pylint: disable=broad-except,broad-exception-caught + log.error(e) + context["graph_result"] = [] + return context - max_deep = self.wk_input.max_deep or 2 - max_graph_items = ( - self.wk_input.max_graph_items or huge_settings.max_graph_items - ) - max_v_prop_len = self.wk_input.max_v_prop_len or 2048 - max_e_prop_len = self.wk_input.max_e_prop_len or 256 - prop_to_match = self.wk_input.prop_to_match - num_gremlin_generate_example = self.wk_input.gremlin_tmpl_num or -1 - gremlin_prompt = ( - self.wk_input.gremlin_prompt or prompt.gremlin_generate_prompt + def _limit_property_query( + self, value: Optional[str], item_type: str + ) -> Optional[str]: + # NOTE: we skip the filter for list/set type (e.g., list of string, add it if needed) + if not self._limit_property or not isinstance(value, str): + return value + + max_len = self._max_v_prop_len if item_type == "v" else self._max_e_prop_len + return value[:max_len] if value else value + + def _process_vertex( + self, + item: Any, + flat_rel: str, + node_cache: Set[str], + prior_edge_str_len: int, + depth: int, + nodes_with_degree: List[str], + use_id_to_match: bool, + v_cache: Set[str], + ) -> Tuple[str, int, int]: + matched_str = ( + item["id"] if use_id_to_match else item["props"][self._prop_to_match] + ) + if matched_str in node_cache: + flat_rel = flat_rel[:-prior_edge_str_len] + return flat_rel, prior_edge_str_len, depth + + node_cache.add(matched_str) + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'v')}" + for k, v in item["props"].items() + if v + ) + + # TODO: we may remove label id or replace with label name + if matched_str in v_cache: + node_str = matched_str + else: + v_cache.add(matched_str) + node_str = f"{item['id']}{{{props_str}}}" + + flat_rel += node_str + nodes_with_degree.append(node_str) + depth += 1 + return flat_rel, prior_edge_str_len, depth + + def _process_edge( + self, + item: Any, + path_str: str, + raw_flat_rel: List[Any], + i: int, + use_id_to_match: bool, + e_cache: Set[Tuple[str, str, str]], + ) -> Tuple[str, int]: + props_str = ", ".join( + f"{k}: {self._limit_property_query(v, 'e')}" + for k, v in item["props"].items() + if v + ) + props_str = f"{{{props_str}}}" if props_str else "" + prev_matched_str = ( + raw_flat_rel[i - 1]["id"] + if use_id_to_match + else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] + ) + + edge_key = (item["inV"], item["label"], item["outV"]) + if edge_key not in e_cache: + e_cache.add(edge_key) + edge_label = f"{item['label']}{props_str}" + else: + edge_label = item["label"] + + edge_str = ( + f"--[{edge_label}]-->" + if item["outV"] == prev_matched_str + else f"<--[{edge_label}]--" + ) + path_str += edge_str + prior_edge_str_len = len(edge_str) + return path_str, prior_edge_str_len + + def _process_path( + self, + path: Any, + use_id_to_match: bool, + v_cache: Set[str], + e_cache: Set[Tuple[str, str, str]], + ) -> Tuple[str, List[str]]: + flat_rel = "" + raw_flat_rel = path["objects"] + assert len(raw_flat_rel) % 2 == 1, "The length of raw_flat_rel should be odd." + + node_cache = set() + prior_edge_str_len = 0 + depth = 0 + nodes_with_degree = [] + + for i, item in enumerate(raw_flat_rel): + if i % 2 == 0: + # Process each vertex + flat_rel, prior_edge_str_len, depth = self._process_vertex( + item, + flat_rel, + node_cache, + prior_edge_str_len, + depth, + nodes_with_degree, + use_id_to_match, + v_cache, + ) + else: + # Process each edge + flat_rel, prior_edge_str_len = self._process_edge( + item, flat_rel, raw_flat_rel, i, use_id_to_match, e_cache + ) + + return flat_rel, nodes_with_degree + + def _update_vertex_degree_list( + self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str] + ) -> None: + for depth, node_str in enumerate(nodes_with_degree): + if depth >= len(vertex_degree_list): + vertex_degree_list.append(set()) + vertex_degree_list[depth].add(node_str) + + def _format_graph_query_result( + self, query_paths + ) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: + use_id_to_match = self._prop_to_match is None + subgraph = set() + subgraph_with_degree = {} + vertex_degree_list: List[Set[str]] = [] + v_cache: Set[str] = set() + e_cache: Set[Tuple[str, str, str]] = set() + + for path in query_paths: + # 1. Process each path + path_str, vertex_with_degree = self._process_path( + path, use_id_to_match, v_cache, e_cache ) + subgraph.add(path_str) + subgraph_with_degree[path_str] = vertex_with_degree + # 2. Update vertex degree list + self._update_vertex_degree_list(vertex_degree_list, vertex_with_degree) + + return subgraph, vertex_degree_list, subgraph_with_degree + + def _get_graph_schema(self, refresh: bool = False) -> str: + if self._schema and not refresh: + return self._schema + + schema = self._client.schema() + vertex_schema = schema.getVertexLabels() + edge_schema = schema.getEdgeLabels() + relationships = schema.getRelations() + + self._schema = ( + f"Vertex properties: {vertex_schema}\n" + f"Edge properties: {edge_schema}\n" + f"Relationships: {relationships}\n" + ) + log.debug("Link(Relation): %s", relationships) + return self._schema + + @staticmethod + def _extract_label_names( + source: str, head: str = "name: ", tail: str = ", " + ) -> List[str]: + result = [] + for s in source.split(head): + end = s.find(tail) + label = s[:end] + if label: + result.append(label) + return result + + def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: + schema = self._get_graph_schema() + vertex_props_str, edge_props_str = schema.split("\n")[:2] + # TODO: rename to vertex (also need update in the schema) + vertex_props_str = ( + vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") + ) + edge_props_str = ( + edge_props_str[len("Edge properties: ") :].strip("[").strip("]") + ) + vertex_labels = self._extract_label_names(vertex_props_str) + edge_labels = self._extract_label_names(edge_props_str) + return vertex_labels, edge_labels + + def _format_graph_from_vertex(self, query_result: List[Any]) -> Set[str]: + knowledge = set() + for item in query_result: + props_str = ", ".join(f"{k}: {v}" for k, v in item["properties"].items()) + node_str = f"{item['id']}{{{props_str}}}" + knowledge.add(node_str) + return knowledge - # Initialize GraphRAGQuery operator - self.graph_rag_query = GraphRAGQuery( - max_deep=max_deep, - max_graph_items=max_graph_items, - max_v_prop_len=max_v_prop_len, - max_e_prop_len=max_e_prop_len, - prop_to_match=prop_to_match, - num_gremlin_generate_example=num_gremlin_generate_example, - gremlin_prompt=gremlin_prompt, + def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: + # 1. Extract params from context + matched_vids = context.get("match_vids") + if isinstance(context.get("max_deep"), int): + self._max_deep = context["max_deep"] + if isinstance(context.get("max_items"), int): + self._max_items = context["max_items"] + if isinstance(context.get("prop_to_match"), str): + self._prop_to_match = context["prop_to_match"] + + # 2. Extract edge_labels from graph schema + _, edge_labels = self._extract_labels_from_schema() + edge_labels_str = ",".join("'" + label + "'" for label in edge_labels) + # TODO: enhance the limit logic later + edge_limit_amount = len(edge_labels) * huge_settings.edge_limit_pre_label + + use_id_to_match = self._prop_to_match is None + if use_id_to_match: + if not matched_vids: + return context + + gremlin_query = VERTEX_QUERY_TPL.format(keywords=matched_vids) + vertexes = self._client.gremlin().exec(gremlin=gremlin_query)["data"] + log.debug("Vids gremlin query: %s", gremlin_query) + + vertex_knowledge = self._format_graph_from_vertex(query_result=vertexes) + paths: List[Any] = [] + # TODO: use generator or asyncio to speed up the query logic + for matched_vid in matched_vids: + gremlin_query = VID_QUERY_NEIGHBOR_TPL.format( + keywords=f"'{matched_vid}'", + max_deep=self._max_deep, + edge_labels=edge_labels_str, + edge_limit=edge_limit_amount, + max_items=self._max_items, + ) + log.debug("Kneighbor gremlin query: %s", gremlin_query) + paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) + + ( + graph_chain_knowledge, + vertex_degree_list, + knowledge_with_degree, + ) = self._format_graph_query_result(query_paths=paths) + + # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) + if not graph_chain_knowledge: + graph_chain_knowledge.update(vertex_knowledge) + if vertex_degree_list: + vertex_degree_list[0].update(vertex_knowledge) + else: + vertex_degree_list.append(vertex_knowledge) + else: + # WARN: When will the query enter here? + keywords = context.get("keywords") + assert keywords, "No related property(keywords) for graph query." + keywords_str = ",".join("'" + kw + "'" for kw in keywords) + gremlin_query = PROPERTY_QUERY_NEIGHBOR_TPL.format( + prop=self._prop_to_match, + keywords=keywords_str, + edge_labels=edge_labels_str, + edge_limit=edge_limit_amount, + max_deep=self._max_deep, + max_items=self._max_items, + ) + log.warning( + "Unable to find vid, downgraded to property query, please confirm if it meets expectation." ) - return super().node_init() - except Exception as e: - log.error(f"Failed to initialize GraphQueryNode: {e}") + paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)[ + "data" + ] + ( + graph_chain_knowledge, + vertex_degree_list, + knowledge_with_degree, + ) = self._format_graph_query_result(query_paths=paths) - return CStatus(-1, f"GraphQueryNode initialization failed: {e}") + context["graph_result"] = list(graph_chain_knowledge) + if context["graph_result"]: + context["graph_result_flag"] = 0 + context["vertex_degree_list"] = [ + list(vertex_degree) for vertex_degree in vertex_degree_list + ] + context["knowledge_with_degree"] = knowledge_with_degree + context["graph_context_head"] = ( + f"The following are graph knowledge in {self._max_deep} depth, e.g:\n" + "`vertexA--[links]-->vertexB<--[links]--vertexC ...`" + "extracted based on key entities as subject:\n" + ) + return context def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: """ @@ -79,15 +479,31 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: return data_json # Execute the graph query (assuming schema and semantic query have been completed in previous nodes) - graph_result = self.graph_rag_query.run(data_json) - data_json.update(graph_result) + self.init_client(data_json) + + # initial flag: -1 means no result, 0 means subgraph query, 1 means gremlin query + data_json["graph_result_flag"] = -1 + # 1. Try to perform a query based on the generated gremlin + if self._num_gremlin_generate_example >= 0: + data_json = self._gremlin_generate_query(data_json) + # 2. Try to perform a query based on subgraph-search if the previous query failed + if not data_json.get("graph_result"): + data_json = self._subgraph_query(data_json) + + if data_json.get("graph_result"): + log.debug( + "Knowledge from Graph:\n%s", "\n".join(data_json["graph_result"]) + ) + else: + log.debug("No Knowledge Extracted from Graph") log.info( - f"Graph query completed, found {len(data_json.get('graph_result', []))} results" + "Graph query completed, found %d results", + len(data_json.get("graph_result", [])), ) return data_json - except Exception as e: - log.error(f"Graph query failed: {e}") + except Exception as e: # pylint: disable=broad-except,broad-exception-caught + log.error("Graph query failed: %s", e) return data_json 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 3face9d63..26d74c5d9 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/hugegraph_node/schema.py @@ -15,6 +15,7 @@ import json +from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.common_op.check_schema import CheckSchema from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager @@ -38,23 +39,23 @@ def _import_schema( ): if from_hugegraph: return SchemaManager(from_hugegraph) - elif from_user_defined: + if from_user_defined: return CheckSchema(from_user_defined) - elif from_extraction: + if from_extraction: raise NotImplementedError("Not implemented yet") - else: - raise ValueError("No input data / invalid schema type") + raise ValueError("No input data / invalid schema type") def node_init(self): - self.schema = self.wk_input.schema - self.schema = self.schema.strip() + if self.wk_input.schema is None: + return CStatus(-1, "Schema message is required in SchemaNode") + self.schema = self.wk_input.schema.strip() if self.schema.startswith("{"): try: schema = json.loads(self.schema) self.check_schema = self._import_schema(from_user_defined=schema) except json.JSONDecodeError as exc: log.error("Invalid JSON format in schema. Please check it again.") - raise ValueError("Invalid JSON format in schema.") from exc + 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) @@ -63,11 +64,6 @@ def node_init(self): def operator_schedule(self, data_json): log.debug("SchemaNode input state: %s", data_json) if self.schema.startswith("{"): - try: - return self.check_schema.run(data_json) - except json.JSONDecodeError as exc: - log.error("Invalid JSON format in schema. Please check it again.") - raise ValueError("Invalid JSON format in schema.") from exc - else: - log.info("Get schema '%s' from graphdb.", self.schema) - return self.schema_manager.run(data_json) + return self.check_schema.run(data_json) + log.info("Get schema '%s' from graphdb.", self.schema) + return self.schema_manager.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py new file mode 100644 index 000000000..8772959d7 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py @@ -0,0 +1,43 @@ +# 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 PyCGraph import CStatus + +from hugegraph_llm.config import llm_settings +from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.nodes.base_node import BaseNode +from hugegraph_llm.operators.index_op.build_gremlin_example_index import ( + BuildGremlinExampleIndex, +) +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +class BuildGremlinExampleIndexNode(BaseNode): + build_gremlin_example_index_op: BuildGremlinExampleIndex + context: WkFlowState = None + wk_input: WkFlowInput = None + + def node_init(self): + if not self.wk_input.examples: + return CStatus(-1, "examples is required in BuildGremlinExampleIndexNode") + examples = self.wk_input.examples + + self.build_gremlin_example_index_op = BuildGremlinExampleIndex( + get_embedding(llm_settings), examples + ) + return super().node_init() + + def operator_schedule(self, data_json): + return self.build_gremlin_example_index_op.run(data_json) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py index bf605aa49..68d2b72f2 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -13,8 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from PyCGraph import CStatus from typing import Dict, Any +from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.semantic_id_query import SemanticIdQuery from hugegraph_llm.models.embeddings.init_embedding import get_embedding @@ -33,59 +33,61 @@ def node_init(self): """ Initialize the semantic ID query operator. """ - try: - graph_name = huge_settings.graph_name - if not graph_name: - return CStatus(-1, "graph_name is required in wk_input") - - embedding = get_embedding(llm_settings) - by = self.wk_input.semantic_by or "keywords" - topk_per_keyword = ( - self.wk_input.topk_per_keyword or huge_settings.topk_per_keyword - ) - topk_per_query = self.wk_input.topk_per_query or 10 - vector_dis_threshold = ( - self.wk_input.vector_dis_threshold or huge_settings.vector_dis_threshold - ) + graph_name = huge_settings.graph_name + if not graph_name: + return CStatus(-1, "graph_name is required in wk_input") - # Initialize the semantic ID query operator - self.semantic_id_query = SemanticIdQuery( - embedding=embedding, - by=by, - topk_per_keyword=topk_per_keyword, - topk_per_query=topk_per_query, - vector_dis_threshold=vector_dis_threshold, - ) + embedding = get_embedding(llm_settings) + by = ( + self.wk_input.semantic_by + if self.wk_input.semantic_by is not None + else "keywords" + ) + topk_per_keyword = ( + self.wk_input.topk_per_keyword + if self.wk_input.topk_per_keyword is not None + else huge_settings.topk_per_keyword + ) + topk_per_query = ( + self.wk_input.topk_per_query + if self.wk_input.topk_per_query is not None + else 10 + ) + vector_dis_threshold = ( + self.wk_input.vector_dis_threshold + if self.wk_input.vector_dis_threshold is not None + else huge_settings.vector_dis_threshold + ) - return super().node_init() - except Exception as e: - log.error(f"Failed to initialize SemanticIdQueryNode: {e}") + # Initialize the semantic ID query operator + self.semantic_id_query = SemanticIdQuery( + embedding=embedding, + by=by, + topk_per_keyword=topk_per_keyword, + topk_per_query=topk_per_query, + vector_dis_threshold=vector_dis_threshold, + ) - return CStatus(-1, f"SemanticIdQueryNode initialization failed: {e}") + return super().node_init() def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: """ Execute the semantic ID query operation. """ - try: - # Get the query text and keywords from input - query = data_json.get("query", "") - keywords = data_json.get("keywords", []) + # Get the query text and keywords from input + query = data_json.get("query", "") + keywords = data_json.get("keywords", []) - if not query and not keywords: - log.warning("No query text or keywords provided for semantic query") - return data_json - - # Perform the semantic query - semantic_result = self.semantic_id_query.run(data_json) + if not query and not keywords: + log.warning("No query text or keywords provided for semantic query") + return data_json - match_vids = semantic_result.get("match_vids", []) - log.info( - f"Semantic query completed, found {len(match_vids)} matching vertex IDs" - ) + # Perform the semantic query + semantic_result = self.semantic_id_query.run(data_json) - return semantic_result + match_vids = semantic_result.get("match_vids", []) + log.info( + "Semantic query completed, found %d matching vertex IDs", len(match_vids) + ) - except Exception as e: - log.error(f"Semantic query failed: {e}") - return data_json + return semantic_result diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py index 48b50acf3..9c8104c6e 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py @@ -32,20 +32,14 @@ def node_init(self): """ Initialize the vector query operator """ - try: - # 从 wk_input 中读取用户配置参数 - embedding = get_embedding(llm_settings) - max_items = ( - self.wk_input.max_items if self.wk_input.max_items is not None else 3 - ) - - self.operator = VectorIndexQuery(embedding=embedding, topk=max_items) - return super().node_init() - except Exception as e: - log.error(f"Failed to initialize VectorQueryNode: {e}") - from PyCGraph import CStatus + # 从 wk_input 中读取用户配置参数 + embedding = get_embedding(llm_settings) + max_items = ( + self.wk_input.max_items if self.wk_input.max_items is not None else 3 + ) - return CStatus(-1, f"VectorQueryNode initialization failed: {e}") + self.operator = VectorIndexQuery(embedding=embedding, topk=max_items) + return super().node_init() def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: """ @@ -64,11 +58,12 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: # Update the state data_json.update(result) log.info( - f"Vector query completed, found {len(result.get('vector_result', []))} results" + "Vector query completed, found %d results", + len(result.get("vector_result", [])), ) return data_json - except Exception as e: - log.error(f"Vector query failed: {e}") + except ValueError as e: + log.error("Vector query failed: %s", e) return data_json diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py index 22b970b4a..6997cd781 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/answer_synthesize_node.py @@ -30,70 +30,57 @@ def node_init(self): """ Initialize the answer synthesis operator. """ - try: - prompt_template = self.wk_input.answer_prompt - raw_answer = self.wk_input.raw_answer or False - vector_only_answer = self.wk_input.vector_only_answer or False - graph_only_answer = self.wk_input.graph_only_answer or False - graph_vector_answer = self.wk_input.graph_vector_answer or False + prompt_template = self.wk_input.answer_prompt + raw_answer = self.wk_input.raw_answer or False + vector_only_answer = self.wk_input.vector_only_answer or False + graph_only_answer = self.wk_input.graph_only_answer or False + graph_vector_answer = self.wk_input.graph_vector_answer or False - self.operator = AnswerSynthesize( - prompt_template=prompt_template, - raw_answer=raw_answer, - vector_only_answer=vector_only_answer, - graph_only_answer=graph_only_answer, - graph_vector_answer=graph_vector_answer, - ) - return super().node_init() - except Exception as e: - log.error(f"Failed to initialize AnswerSynthesizeNode: {e}") - from PyCGraph import CStatus - - return CStatus(-1, f"AnswerSynthesizeNode initialization failed: {e}") + self.operator = AnswerSynthesize( + prompt_template=prompt_template, + raw_answer=raw_answer, + vector_only_answer=vector_only_answer, + graph_only_answer=graph_only_answer, + graph_vector_answer=graph_vector_answer, + ) + return super().node_init() def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: """ Execute the answer synthesis operation. """ - try: - if self.getGParamWithNoEmpty("wkflow_input").stream: - # Streaming mode: return a generator for streaming output - data_json["stream_generator"] = self.operator.run_streaming(data_json) - return data_json - else: - # Non-streaming mode: execute answer synthesis - result = self.operator.run(data_json) - - # Record the types of answers generated - answer_types = [] - if result.get("raw_answer"): - answer_types.append("raw") - if result.get("vector_only_answer"): - answer_types.append("vector_only") - if result.get("graph_only_answer"): - answer_types.append("graph_only") - if result.get("graph_vector_answer"): - answer_types.append("graph_vector") + if self.getGParamWithNoEmpty("wkflow_input").stream: + # Streaming mode: return a generator for streaming output + data_json["stream_generator"] = self.operator.run_streaming(data_json) + return data_json + # Non-streaming mode: execute answer synthesis + result = self.operator.run(data_json) - log.info( - f"Answer synthesis completed for types: {', '.join(answer_types)}" - ) + # Record the types of answers generated + answer_types = [] + if result.get("raw_answer"): + answer_types.append("raw") + if result.get("vector_only_answer"): + answer_types.append("vector_only") + if result.get("graph_only_answer"): + answer_types.append("graph_only") + if result.get("graph_vector_answer"): + answer_types.append("graph_vector") - # Print enabled answer types according to self.wk_input configuration - wk_input_types = [] - if getattr(self.wk_input, "raw_answer", False): - wk_input_types.append("raw") - if getattr(self.wk_input, "vector_only_answer", False): - wk_input_types.append("vector_only") - if getattr(self.wk_input, "graph_only_answer", False): - wk_input_types.append("graph_only") - if getattr(self.wk_input, "graph_vector_answer", False): - wk_input_types.append("graph_vector") - log.info( - f"Enabled answer types according to wk_input config: {', '.join(wk_input_types)}" - ) - return result + log.info("Answer synthesis completed for types: %s", ", ".join(answer_types)) - except Exception as e: - log.error(f"Answer synthesis failed: {e}") - return data_json + # Print enabled answer types according to self.wk_input configuration + wk_input_types = [] + if getattr(self.wk_input, "raw_answer", False): + wk_input_types.append("raw") + if getattr(self.wk_input, "vector_only_answer", False): + wk_input_types.append("vector_only") + if getattr(self.wk_input, "graph_only_answer", False): + wk_input_types.append("graph_only") + if getattr(self.wk_input, "graph_vector_answer", False): + wk_input_types.append("graph_vector") + log.info( + "Enabled answer types according to wk_input config: %s", + ", ".join(wk_input_types), + ) + return result 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 628765f58..3c9bf2308 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 @@ -48,5 +48,6 @@ def node_init(self): def operator_schedule(self, data_json): if self.extract_type == "triples": return self.info_extract.run(data_json) - elif self.extract_type == "property_graph": + if self.extract_type == "property_graph": return self.property_graph_extract.run(data_json) + raise ValueError("Unsupport extract type") diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py index 76fc06eb3..60542ddc1 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/keyword_extract_node.py @@ -14,7 +14,6 @@ # limitations under the License. from typing import Dict, Any -from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.llm_op.keyword_extract import KeywordExtract @@ -32,29 +31,17 @@ def node_init(self): """ Initialize the keyword extraction operator. """ - try: - max_keywords = ( - self.wk_input.max_keywords - if self.wk_input.max_keywords is not None - else 5 - ) - language = ( - self.wk_input.language - if self.wk_input.language is not None - else "english" - ) - extract_template = self.wk_input.keywords_extract_prompt + max_keywords = ( + self.wk_input.max_keywords if self.wk_input.max_keywords is not None else 5 + ) + extract_template = self.wk_input.keywords_extract_prompt - self.operator = KeywordExtract( - text=self.wk_input.query, - max_keywords=max_keywords, - language=language, - extract_template=extract_template, - ) - return super().node_init() - except Exception as e: - log.error(f"Failed to initialize KeywordExtractNode: {e}") - return CStatus(-1, f"KeywordExtractNode initialization failed: {e}") + self.operator = KeywordExtract( + text=self.wk_input.query, + max_keywords=max_keywords, + extract_template=extract_template, + ) + return super().node_init() def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: """ @@ -67,12 +54,12 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: log.warning("Keyword extraction result missing 'keywords' field") result["keywords"] = [] - log.info(f"Extracted keywords: {result.get('keywords', [])}") + log.info("Extracted keywords: %s", result.get("keywords", [])) return result - except Exception as e: - log.error(f"Keyword extraction failed: {e}") + except ValueError as e: + log.error("Keyword extraction failed: %s", e) # Add error flag to indicate failure error_result = data_json.copy() error_result["error"] = str(e) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py index 408adb10a..1ef7e5c55 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/schema_build.py @@ -61,12 +61,16 @@ def node_init(self): # few_shot_schema: already parsed dict or raw JSON string few_shot_schema = {} - fss_src = self.wk_input.few_shot_schema if self.wk_input.few_shot_schema else None + fss_src = ( + self.wk_input.few_shot_schema if self.wk_input.few_shot_schema else None + ) if fss_src: try: few_shot_schema = json.loads(fss_src) except json.JSONDecodeError as e: - return CStatus(-1, f"Few Shot Schema is not in a valid JSON format: {e}") + return CStatus( + -1, f"Few Shot Schema is not in a valid JSON format: {e}" + ) _context_payload = { "raw_texts": raw_texts, @@ -82,6 +86,6 @@ def operator_schedule(self, data_json): schema_result = self.schema_builder.run(data_json) return {"schema": schema_result} - except Exception as e: + except (ValueError, RuntimeError) as e: log.error("Failed to generate schema: %s", e) return {"schema": f"Schema generation failed: {e}"} diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py index a36831526..0904b9920 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/llm_node/text2gremlin.py @@ -18,7 +18,6 @@ import json from typing import Any, Dict, Optional -from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize @@ -27,13 +26,14 @@ def _stable_schema_string(state_json: Dict[str, Any]) -> str: - if "simple_schema" in state_json and state_json["simple_schema"] is not None: - return json.dumps( - state_json["simple_schema"], ensure_ascii=False, sort_keys=True - ) - if "schema" in state_json and state_json["schema"] is not None: - return json.dumps(state_json["schema"], ensure_ascii=False, sort_keys=True) - return "" + val = state_json.get("simple_schema") + if val is None: + val = state_json.get("schema") + if val is None: + return "" + if isinstance(val, str): + return val + return json.dumps(val, ensure_ascii=False, sort_keys=True) class Text2GremlinNode(BaseNode): @@ -56,7 +56,7 @@ def node_init(self): vertices=None, gremlin_prompt=gremlin_prompt, ) - return CStatus() + return super().node_init() def operator_schedule(self, data_json: Dict[str, Any]): # Ensure query exists in context; return empty if not provided diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/util.py b/hugegraph-llm/src/hugegraph_llm/nodes/util.py index 60bdc2e86..d1ac69657 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/util.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/util.py @@ -13,15 +13,26 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any + from PyCGraph import CStatus -def init_context(obj) -> CStatus: - try: - obj.context = obj.getGParamWithNoEmpty("wkflow_state") - obj.wk_input = obj.getGParamWithNoEmpty("wkflow_input") - if obj.context is None or obj.wk_input is None: - return CStatus(-1, "Required workflow parameters not found") - return CStatus() - except Exception as e: - return CStatus(-1, f"Failed to initialize context: {str(e)}") +def init_context(obj: Any) -> CStatus: + """ + Initialize workflow context for a node. + + Retrieves wkflow_state and wkflow_input from obj's global parameters + and assigns them to obj.context and obj.wk_input respectively. + + Args: + obj: Node object with getGParamWithNoEmpty method + + Returns: + CStatus: Empty status on success, error status with code -1 on failure + """ + obj.context = obj.getGParamWithNoEmpty("wkflow_state") + obj.wk_input = obj.getGParamWithNoEmpty("wkflow_input") + if obj.context is None or obj.wk_input is None: + return CStatus(-1, "Required workflow parameters not found") + return CStatus() diff --git a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py b/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py deleted file mode 100644 index 70f3d27d2..000000000 --- a/hugegraph-llm/src/hugegraph_llm/operators/gremlin_generate_task.py +++ /dev/null @@ -1,81 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -from typing import Optional, List - -from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.models.llms.base import BaseLLM -from hugegraph_llm.operators.common_op.check_schema import CheckSchema -from hugegraph_llm.operators.common_op.print_result import PrintResult -from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager -from hugegraph_llm.operators.index_op.build_gremlin_example_index import BuildGremlinExampleIndex -from hugegraph_llm.operators.index_op.gremlin_example_index_query import GremlinExampleIndexQuery -from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize -from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm - - -class GremlinGenerator: - def __init__(self, llm: BaseLLM, embedding: BaseEmbedding): - self.embedding = [] - self.llm = llm - self.embedding = embedding - self.result = None - self.operators = [] - - def clear(self): - self.operators = [] - return self - - def example_index_build(self, examples): - self.operators.append(BuildGremlinExampleIndex(self.embedding, examples)) - return self - - def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): - if from_hugegraph: - self.operators.append(SchemaManager(from_hugegraph)) - elif from_user_defined: - self.operators.append(CheckSchema(from_user_defined)) - elif from_extraction: - raise NotImplementedError("Not implemented yet") - else: - raise ValueError("No input data / invalid schema type") - return self - - def example_index_query(self, num_examples): - self.operators.append(GremlinExampleIndexQuery(self.embedding, num_examples)) - return self - - def gremlin_generate_synthesize( - self, schema, gremlin_prompt: Optional[str] = None, vertices: Optional[List[str]] = None - ): - self.operators.append(GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt)) - return self - - def print_result(self): - self.operators.append(PrintResult()) - return self - - @log_time("total time") - @record_rpm - def run(self, **kwargs): - context = kwargs - for operator in self.operators: - context = self._run_operator(operator, context) - return context - - @log_operator_time - def _run_operator(self, operator, context): - return operator.run(context) 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 52626b72b..ba4392f7c 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 @@ -40,7 +40,6 @@ def run(self, data: dict) -> Dict[str, Any]: schema = data.get("schema") vertices = data.get("vertices", []) edges = data.get("edges", []) - print(f"get schema {schema}") if not vertices and not edges: log.critical( "(Loading) Both vertices and edges are empty. Please check the input data again." @@ -50,7 +49,9 @@ def run(self, data: dict) -> Dict[str, Any]: if not schema: # TODO: ensure the function works correctly (update the logic later) self.schema_free_mode(data.get("triples", [])) - log.warning("Using schema_free mode, could try schema_define mode for better effect!") + log.warning( + "Using schema_free mode, could try schema_define mode for better effect!" + ) else: self.init_schema_if_need(schema) self.load_into_graph(vertices, edges, schema) @@ -66,7 +67,9 @@ def _set_default_property(self, key, input_properties, property_label_map): # list or set default_value = [] input_properties[key] = default_value - log.warning("Property '%s' missing in vertex, set to '%s' for now", key, default_value) + log.warning( + "Property '%s' missing in vertex, set to '%s' for now", key, default_value + ) def _handle_graph_creation(self, func, *args, **kwargs): try: @@ -80,9 +83,13 @@ 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) - vertex_label_map = {v_label["name"]: v_label for v_label in schema["vertexlabels"]} + 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"]} + property_label_map = { + p_label["name"]: p_label for p_label in schema["propertykeys"] + } for vertex in vertices: input_label = vertex["label"] @@ -98,7 +105,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- vertex_label = vertex_label_map[input_label] primary_keys = vertex_label["primary_keys"] nullable_keys = vertex_label.get("nullable_keys", []) - non_null_keys = [key for key in vertex_label["properties"] if key not in nullable_keys] + non_null_keys = [ + key for key in vertex_label["properties"] if key not in nullable_keys + ] has_problem = False # 2. Handle primary-keys mode vertex @@ -130,7 +139,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- # 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) + 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(): @@ -167,7 +178,9 @@ def load_into_graph(self, vertices, edges, schema): # pylint: disable=too-many- 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) + self._handle_graph_creation( + self.client.graph().addEdge, label, start, end, properties + ) def init_schema_if_need(self, schema: dict): properties = schema["propertykeys"] @@ -191,18 +204,20 @@ def init_schema_if_need(self, schema: dict): source_vertex_label = edge["source_label"] target_vertex_label = edge["target_label"] properties = edge["properties"] - self.schema.edgeLabel(edge_label).sourceLabel(source_vertex_label).targetLabel( - target_vertex_label - ).properties(*properties).nullableKeys(*properties).ifNotExist().create() + self.schema.edgeLabel(edge_label).sourceLabel( + source_vertex_label + ).targetLabel(target_vertex_label).properties(*properties).nullableKeys( + *properties + ).ifNotExist().create() def schema_free_mode(self, 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( - "name" - ).ifNotExist().create() + self.schema.edgeLabel("edge").sourceLabel("vertex").targetLabel( + "vertex" + ).properties("name").ifNotExist().create() self.schema.indexLabel("vertexByName").onV("vertex").by( "name" @@ -262,7 +277,9 @@ def _set_property_data_type(self, property_key, data_type): log.warning("UUID type is not supported, use text instead") property_key.asText() else: - log.error("Unknown data type %s for property_key %s", data_type, property_key) + log.error( + "Unknown data type %s for property_key %s", data_type, property_key + ) def _set_property_cardinality(self, property_key, cardinality): if cardinality == PropertyCardinality.SINGLE: @@ -272,9 +289,13 @@ def _set_property_cardinality(self, property_key, cardinality): elif cardinality == PropertyCardinality.SET: property_key.valueSet() else: - log.error("Unknown cardinality %s for property_key %s", cardinality, property_key) + 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: + def _check_property_data_type( + self, data_type: str, cardinality: str, value + ) -> bool: if cardinality in ( PropertyCardinality.LIST.value, PropertyCardinality.SET.value, @@ -304,7 +325,9 @@ def _check_single_data_type(self, data_type: str, value) -> bool: if data_type in (PropertyDataType.TEXT.value, PropertyDataType.UUID.value): return isinstance(value, str) # TODO: check ok below - if data_type == PropertyDataType.DATE.value: # the format should be "yyyy-MM-dd" + if ( + data_type == PropertyDataType.DATE.value + ): # the format should be "yyyy-MM-dd" import re return isinstance(value, str) and re.match(r"^\d{4}-\d{2}-\d{2}$", value) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py b/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py deleted file mode 100644 index bcff5f07b..000000000 --- a/hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py +++ /dev/null @@ -1,455 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -import json -from typing import Any, Dict, Optional, List, Set, Tuple - -from hugegraph_llm.config import huge_settings, prompt -from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.models.llms.base import BaseLLM -from hugegraph_llm.operators.gremlin_generate_task import GremlinGenerator -from hugegraph_llm.utils.log import log -from pyhugegraph.client import PyHugeClient - -# TODO: remove 'as('subj)' step -VERTEX_QUERY_TPL = "g.V({keywords}).limit(8).as('subj').toList()" - -# TODO: we could use a simpler query (like kneighbor-api to get the edges) -# TODO: test with profile()/explain() to speed up the query -VID_QUERY_NEIGHBOR_TPL = """\ -g.V({keywords}) -.repeat( - bothE({edge_labels}).limit({edge_limit}).otherV().dedup() -).times({max_deep}).emit() -.simplePath() -.path() -.by(project('label', 'id', 'props') - .by(label()) - .by(id()) - .by(valueMap().by(unfold())) -) -.by(project('label', 'inV', 'outV', 'props') - .by(label()) - .by(inV().id()) - .by(outV().id()) - .by(valueMap().by(unfold())) -) -.limit({max_items}) -.toList() -""" - -PROPERTY_QUERY_NEIGHBOR_TPL = """\ -g.V().has('{prop}', within({keywords})) -.repeat( - bothE({edge_labels}).limit({edge_limit}).otherV().dedup() -).times({max_deep}).emit() -.simplePath() -.path() -.by(project('label', 'props') - .by(label()) - .by(valueMap().by(unfold())) -) -.by(project('label', 'inV', 'outV', 'props') - .by(label()) - .by(inV().values('{prop}')) - .by(outV().values('{prop}')) - .by(valueMap().by(unfold())) -) -.limit({max_items}) -.toList() -""" - - -class GraphRAGQuery: - def __init__( - self, - max_deep: int = 2, - max_graph_items: int = huge_settings.max_graph_items, - prop_to_match: Optional[str] = None, - llm: Optional[BaseLLM] = None, - embedding: Optional[BaseEmbedding] = None, - max_v_prop_len: Optional[int] = 2048, - max_e_prop_len: Optional[int] = 256, - num_gremlin_generate_example: Optional[int] = -1, - gremlin_prompt: Optional[str] = None, - ): - 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, - ) - self._max_deep = max_deep - self._max_items = max_graph_items - self._prop_to_match = prop_to_match - self._schema = "" - self._limit_property = huge_settings.limit_property.lower() == "true" - self._max_v_prop_len = max_v_prop_len - self._max_e_prop_len = max_e_prop_len - self._gremlin_generator = GremlinGenerator( - llm=llm, - embedding=embedding, - ) - self._num_gremlin_generate_example = num_gremlin_generate_example - self._gremlin_prompt = gremlin_prompt or prompt.gremlin_generate_prompt - - def run(self, context: Dict[str, Any]) -> Dict[str, Any]: - self.init_client(context) - - # initial flag: -1 means no result, 0 means subgraph query, 1 means gremlin query - context["graph_result_flag"] = -1 - # 1. Try to perform a query based on the generated gremlin - if self._num_gremlin_generate_example >= 0: - context = self._gremlin_generate_query(context) - # 2. Try to perform a query based on subgraph-search if the previous query failed - if not context.get("graph_result"): - context = self._subgraph_query(context) - - if context.get("graph_result"): - log.debug("Knowledge from Graph:\n%s", "\n".join(context["graph_result"])) - else: - log.debug("No Knowledge Extracted from Graph") - return context - - def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]: - query = context["query"] - vertices = context.get("match_vids") - query_embedding = context.get("query_embedding") - - self._gremlin_generator.clear() - self._gremlin_generator.example_index_query(num_examples=self._num_gremlin_generate_example) - gremlin_response = self._gremlin_generator.gremlin_generate_synthesize( - context["simple_schema"], vertices=vertices, gremlin_prompt=self._gremlin_prompt - ).run(query=query, query_embedding=query_embedding) - if self._num_gremlin_generate_example > 0: - gremlin = gremlin_response["result"] - else: - gremlin = gremlin_response["raw_result"] - log.info("Generated gremlin: %s", gremlin) - context["gremlin"] = gremlin - try: - result = self._client.gremlin().exec(gremlin=gremlin)["data"] - if result == [None]: - result = [] - context["graph_result"] = [json.dumps(item, ensure_ascii=False) for item in result] - if context["graph_result"]: - context["graph_result_flag"] = 1 - context["graph_context_head"] = ( - f"The following are graph query result " f"from gremlin query `{gremlin}`.\n" - ) - except Exception as e: # pylint: disable=broad-except - log.error(e) - context["graph_result"] = "" - return context - - def _subgraph_query(self, context: Dict[str, Any]) -> Dict[str, Any]: - # 1. Extract params from context - matched_vids = context.get("match_vids") - if isinstance(context.get("max_deep"), int): - self._max_deep = context["max_deep"] - if isinstance(context.get("max_items"), int): - self._max_items = context["max_items"] - if isinstance(context.get("prop_to_match"), str): - self._prop_to_match = context["prop_to_match"] - - # 2. Extract edge_labels from graph schema - _, edge_labels = self._extract_labels_from_schema() - edge_labels_str = ",".join("'" + label + "'" for label in edge_labels) - # TODO: enhance the limit logic later - edge_limit_amount = len(edge_labels) * huge_settings.edge_limit_pre_label - - use_id_to_match = self._prop_to_match is None - if use_id_to_match: - if not matched_vids: - return context - - gremlin_query = VERTEX_QUERY_TPL.format(keywords=matched_vids) - vertexes = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - log.debug("Vids gremlin query: %s", gremlin_query) - - vertex_knowledge = self._format_graph_from_vertex(query_result=vertexes) - paths: List[Any] = [] - # TODO: use generator or asyncio to speed up the query logic - for matched_vid in matched_vids: - gremlin_query = VID_QUERY_NEIGHBOR_TPL.format( - keywords=f"'{matched_vid}'", - max_deep=self._max_deep, - edge_labels=edge_labels_str, - edge_limit=edge_limit_amount, - max_items=self._max_items, - ) - log.debug("Kneighbor gremlin query: %s", gremlin_query) - paths.extend(self._client.gremlin().exec(gremlin=gremlin_query)["data"]) - - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( - self._format_graph_query_result(query_paths=paths) - ) - - # TODO: we may need to optimize the logic here with global deduplication (may lack some single vertex) - if not graph_chain_knowledge: - graph_chain_knowledge.update(vertex_knowledge) - if vertex_degree_list: - vertex_degree_list[0].update(vertex_knowledge) - else: - vertex_degree_list.append(vertex_knowledge) - else: - # WARN: When will the query enter here? - keywords = context.get("keywords") - assert keywords, "No related property(keywords) for graph query." - keywords_str = ",".join("'" + kw + "'" for kw in keywords) - gremlin_query = PROPERTY_QUERY_NEIGHBOR_TPL.format( - prop=self._prop_to_match, - keywords=keywords_str, - edge_labels=edge_labels_str, - edge_limit=edge_limit_amount, - max_deep=self._max_deep, - max_items=self._max_items, - ) - log.warning( - "Unable to find vid, downgraded to property query, please confirm if it meets expectation." - ) - - paths: List[Any] = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - graph_chain_knowledge, vertex_degree_list, knowledge_with_degree = ( - self._format_graph_query_result(query_paths=paths) - ) - - context["graph_result"] = list(graph_chain_knowledge) - if context["graph_result"]: - context["graph_result_flag"] = 0 - context["vertex_degree_list"] = [ - list(vertex_degree) for vertex_degree in vertex_degree_list - ] - context["knowledge_with_degree"] = knowledge_with_degree - context["graph_context_head"] = ( - f"The following are graph knowledge in {self._max_deep} depth, e.g:\n" - "`vertexA--[links]-->vertexB<--[links]--vertexC ...`" - "extracted based on key entities as subject:\n" - ) - return context - - # TODO: move this method to a util file for reuse (remove self param) - def init_client(self, context): - """Initialize the HugeGraph client from context or default settings.""" - # pylint: disable=R0915 (too-many-statements) - if self._client is None: - if isinstance(context.get("graph_client"), PyHugeClient): - self._client = context["graph_client"] - else: - url = context.get("url") or "http://localhost:8080" - graph = context.get("graph") or "hugegraph" - user = context.get("user") or "admin" - pwd = context.get("pwd") or "admin" - gs = context.get("graphspace") or None - self._client = PyHugeClient(url, graph, user, pwd, gs) - assert self._client is not None, "No valid graph to search." - - def get_vertex_details(self, vertex_ids: List[str]) -> List[Dict[str, Any]]: - if not vertex_ids: - return [] - - formatted_ids = ", ".join(f"'{vid}'" for vid in vertex_ids) - gremlin_query = f"g.V({formatted_ids}).limit(20)" - result = self._client.gremlin().exec(gremlin=gremlin_query)["data"] - return result - - def _format_graph_from_vertex(self, query_result: List[Any]) -> Set[str]: - knowledge = set() - for item in query_result: - props_str = ", ".join(f"{k}: {v}" for k, v in item["properties"].items()) - node_str = f"{item['id']}{{{props_str}}}" - knowledge.add(node_str) - return knowledge - - def _format_graph_query_result( - self, query_paths - ) -> Tuple[Set[str], List[Set[str]], Dict[str, List[str]]]: - use_id_to_match = self._prop_to_match is None - subgraph = set() - subgraph_with_degree = {} - vertex_degree_list: List[Set[str]] = [] - v_cache: Set[str] = set() - e_cache: Set[Tuple[str, str, str]] = set() - - for path in query_paths: - # 1. Process each path - path_str, vertex_with_degree = self._process_path( - path, use_id_to_match, v_cache, e_cache - ) - subgraph.add(path_str) - subgraph_with_degree[path_str] = vertex_with_degree - # 2. Update vertex degree list - self._update_vertex_degree_list(vertex_degree_list, vertex_with_degree) - - return subgraph, vertex_degree_list, subgraph_with_degree - - def _process_path( - self, - path: Any, - use_id_to_match: bool, - v_cache: Set[str], - e_cache: Set[Tuple[str, str, str]], - ) -> Tuple[str, List[str]]: - flat_rel = "" - raw_flat_rel = path["objects"] - assert len(raw_flat_rel) % 2 == 1, "The length of raw_flat_rel should be odd." - - node_cache = set() - prior_edge_str_len = 0 - depth = 0 - nodes_with_degree = [] - - for i, item in enumerate(raw_flat_rel): - if i % 2 == 0: - # Process each vertex - flat_rel, prior_edge_str_len, depth = self._process_vertex( - item, - flat_rel, - node_cache, - prior_edge_str_len, - depth, - nodes_with_degree, - use_id_to_match, - v_cache, - ) - else: - # Process each edge - flat_rel, prior_edge_str_len = self._process_edge( - item, flat_rel, raw_flat_rel, i, use_id_to_match, e_cache - ) - - return flat_rel, nodes_with_degree - - def _process_vertex( - self, - item: Any, - flat_rel: str, - node_cache: Set[str], - prior_edge_str_len: int, - depth: int, - nodes_with_degree: List[str], - use_id_to_match: bool, - v_cache: Set[str], - ) -> Tuple[str, int, int]: - matched_str = item["id"] if use_id_to_match else item["props"][self._prop_to_match] - if matched_str in node_cache: - flat_rel = flat_rel[:-prior_edge_str_len] - return flat_rel, prior_edge_str_len, depth - - node_cache.add(matched_str) - props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'v')}" for k, v in item["props"].items() if v - ) - - # TODO: we may remove label id or replace with label name - if matched_str in v_cache: - node_str = matched_str - else: - v_cache.add(matched_str) - node_str = f"{item['id']}{{{props_str}}}" - - flat_rel += node_str - nodes_with_degree.append(node_str) - depth += 1 - return flat_rel, prior_edge_str_len, depth - - def _process_edge( - self, - item: Any, - path_str: str, - raw_flat_rel: List[Any], - i: int, - use_id_to_match: bool, - e_cache: Set[Tuple[str, str, str]], - ) -> Tuple[str, int]: - props_str = ", ".join( - f"{k}: {self._limit_property_query(v, 'e')}" for k, v in item["props"].items() if v - ) - props_str = f"{{{props_str}}}" if props_str else "" - prev_matched_str = ( - raw_flat_rel[i - 1]["id"] - if use_id_to_match - else (raw_flat_rel)[i - 1]["props"][self._prop_to_match] - ) - - edge_key = (item["inV"], item["label"], item["outV"]) - if edge_key not in e_cache: - e_cache.add(edge_key) - edge_label = f"{item['label']}{props_str}" - else: - edge_label = item["label"] - - edge_str = ( - f"--[{edge_label}]-->" if item["outV"] == prev_matched_str else f"<--[{edge_label}]--" - ) - path_str += edge_str - prior_edge_str_len = len(edge_str) - return path_str, prior_edge_str_len - - def _update_vertex_degree_list( - self, vertex_degree_list: List[Set[str]], nodes_with_degree: List[str] - ) -> None: - for depth, node_str in enumerate(nodes_with_degree): - if depth >= len(vertex_degree_list): - vertex_degree_list.append(set()) - vertex_degree_list[depth].add(node_str) - - def _extract_labels_from_schema(self) -> Tuple[List[str], List[str]]: - schema = self._get_graph_schema() - vertex_props_str, edge_props_str = schema.split("\n")[:2] - # TODO: rename to vertex (also need update in the schema) - vertex_props_str = vertex_props_str[len("Vertex properties: ") :].strip("[").strip("]") - edge_props_str = edge_props_str[len("Edge properties: ") :].strip("[").strip("]") - vertex_labels = self._extract_label_names(vertex_props_str) - edge_labels = self._extract_label_names(edge_props_str) - return vertex_labels, edge_labels - - @staticmethod - def _extract_label_names(source: str, head: str = "name: ", tail: str = ", ") -> List[str]: - result = [] - for s in source.split(head): - end = s.find(tail) - label = s[:end] - if label: - result.append(label) - return result - - def _get_graph_schema(self, refresh: bool = False) -> str: - if self._schema and not refresh: - return self._schema - - schema = self._client.schema() - vertex_schema = schema.getVertexLabels() - edge_schema = schema.getEdgeLabels() - relationships = schema.getRelations() - - self._schema = ( - f"Vertex properties: {vertex_schema}\n" - f"Edge properties: {edge_schema}\n" - f"Relationships: {relationships}\n" - ) - log.debug("Link(Relation): %s", relationships) - return self._schema - - def _limit_property_query(self, value: Optional[str], item_type: str) -> Optional[str]: - # NOTE: we skip the filter for list/set type (e.g., list of string, add it if needed) - if not self._limit_property or not isinstance(value, str): - return value - - max_len = self._max_v_prop_len if item_type == "v" else self._max_e_prop_len - return value[:max_len] if value else value 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 5689a59ac..2ed4e840a 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 @@ -37,11 +37,15 @@ def __init__(self, embedding: BaseEmbedding): self.folder_name = get_index_folder_name( huge_settings.graph_name, huge_settings.graph_space ) - self.index_dir = str(os.path.join(resource_path, self.folder_name, "graph_vids")) + self.index_dir = str( + os.path.join(resource_path, self.folder_name, "graph_vids") + ) self.filename_prefix = get_filename_prefix( llm_settings.embedding_type, getattr(embedding, "model_name", None) ) - self.vid_index = VectorIndex.from_index_file(self.index_dir, self.filename_prefix) + self.vid_index = VectorIndex.from_index_file( + self.index_dir, self.filename_prefix + ) self.embedding = embedding self.sm = SchemaManager(huge_settings.graph_name) @@ -50,19 +54,27 @@ def _extract_names(self, vertices: list[str]) -> list[str]: def run(self, context: Dict[str, Any]) -> Dict[str, Any]: vertexlabels = self.sm.schema.getSchema()["vertexlabels"] - all_pk_flag = all(data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels) + all_pk_flag = bool(vertexlabels) and all( + data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels + ) past_vids = self.vid_index.properties # TODO: We should build vid vector index separately, especially when the vertices may be very large - present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py + present_vids = context[ + "vertices" + ] # Warning: data truncated by fetch_graph_data.py removed_vids = set(past_vids) - set(present_vids) removed_num = self.vid_index.remove(removed_vids) added_vids = list(set(present_vids) - set(past_vids)) if added_vids: - vids_to_process = self._extract_names(added_vids) if all_pk_flag else added_vids - added_embeddings = asyncio.run(get_embeddings_parallel(self.embedding, vids_to_process)) + vids_to_process = ( + self._extract_names(added_vids) if all_pk_flag else added_vids + ) + added_embeddings = asyncio.run( + get_embeddings_parallel(self.embedding, 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.to_index_file(self.index_dir, self.filename_prefix) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py b/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py deleted file mode 100644 index 3b5c63103..000000000 --- a/hugegraph-llm/src/hugegraph_llm/operators/kg_construction_task.py +++ /dev/null @@ -1,120 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - - -from typing import Dict, Any, Optional, Literal, Union, List - -from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.models.llms.base import BaseLLM -from hugegraph_llm.operators.common_op.check_schema import CheckSchema -from hugegraph_llm.operators.common_op.print_result import PrintResult -from hugegraph_llm.operators.document_op.chunk_split import ChunkSplit -from hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph import Commit2Graph -from hugegraph_llm.operators.hugegraph_op.fetch_graph_data import FetchGraphData -from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager -from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex -from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex -from hugegraph_llm.operators.llm_op.disambiguate_data import DisambiguateData -from hugegraph_llm.operators.llm_op.info_extract import InfoExtract -from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract -from hugegraph_llm.operators.llm_op.schema_build import SchemaBuilder -from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm -from pyhugegraph.client import PyHugeClient - - -class KgBuilder: - def __init__( - self, - llm: BaseLLM, - embedding: Optional[BaseEmbedding] = None, - graph: Optional[PyHugeClient] = None, - ): - self.operators = [] - self.llm = llm - self.embedding = embedding - self.graph = graph - self.result = None - - def import_schema(self, from_hugegraph=None, from_extraction=None, from_user_defined=None): - if from_hugegraph: - self.operators.append(SchemaManager(from_hugegraph)) - elif from_user_defined: - self.operators.append(CheckSchema(from_user_defined)) - elif from_extraction: - raise NotImplementedError("Not implemented yet") - else: - raise ValueError("No input data / invalid schema type") - return self - - def fetch_graph_data(self): - self.operators.append(FetchGraphData(self.graph)) - return self - - def chunk_split( - self, - text: Union[str, List[str]], # text to be split - split_type: Literal["document", "paragraph", "sentence"] = "document", - language: Literal["zh", "en"] = "zh", - ): - self.operators.append(ChunkSplit(text, split_type, language)) - return self - - def extract_info( - self, - example_prompt: Optional[str] = None, - extract_type: Literal["triples", "property_graph"] = "triples", - ): - if extract_type == "triples": - self.operators.append(InfoExtract(self.llm, example_prompt)) - elif extract_type == "property_graph": - self.operators.append(PropertyGraphExtract(self.llm, example_prompt)) - return self - - def disambiguate_word_sense(self): - self.operators.append(DisambiguateData(self.llm)) - return self - - def commit_to_hugegraph(self): - self.operators.append(Commit2Graph()) - return self - - def build_vertex_id_semantic_index(self): - self.operators.append(BuildSemanticIndex(self.embedding)) - return self - - def build_vector_index(self): - self.operators.append(BuildVectorIndex(self.embedding)) - return self - - def print_result(self): - self.operators.append(PrintResult()) - return self - - def build_schema(self): - self.operators.append(SchemaBuilder(self.llm)) - return self - - @log_time("total time") - @record_rpm - def run(self, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: - for operator in self.operators: - context = self._run_operator(operator, context) - return context - - @log_operator_time - def _run_operator(self, operator, context) -> Dict[str, Any]: - return operator.run(context) diff --git a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py index 32ed9651e..48369b4ec 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py @@ -22,7 +22,9 @@ from hugegraph_llm.config import prompt, llm_settings from hugegraph_llm.models.llms.base import BaseLLM from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.document_op.textrank_word_extract import MultiLingualTextRank +from hugegraph_llm.operators.document_op.textrank_word_extract import ( + MultiLingualTextRank, +) from hugegraph_llm.utils.log import log KEYWORDS_EXTRACT_TPL = prompt.keywords_extract_prompt @@ -43,8 +45,8 @@ def __init__( self._extract_template = extract_template or KEYWORDS_EXTRACT_TPL self._extract_method = llm_settings.keyword_extract_type.lower() self._textrank_model = MultiLingualTextRank( - keyword_num=max_keywords, - window_size=llm_settings.window_size) + keyword_num=max_keywords, window_size=llm_settings.window_size + ) def run(self, context: Dict[str, Any]) -> Dict[str, Any]: if self._query is None: @@ -66,7 +68,11 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: max_keyword_num = self._max_keywords self._max_keywords = max(1, max_keyword_num) - method = (context.get("extract_method", self._extract_method) or "LLM").strip().lower() + method = ( + (context.get("extract_method", self._extract_method) or "LLM") + .strip() + .lower() + ) if method == "llm": # LLM method ranks = self._extract_with_llm() @@ -82,7 +88,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: keywords = [] if not ranks else sorted(ranks, key=ranks.get, reverse=True) keywords = [k.replace("'", "") for k in keywords] - context["keywords"] = keywords[:self._max_keywords] + context["keywords"] = keywords[: self._max_keywords] log.info("User Query: %s\nKeywords: %s", self._query, context["keywords"]) # extracting keywords & expanding synonyms increase the call count by 1 @@ -101,7 +107,7 @@ def _extract_with_llm(self) -> Dict[str, float]: return keywords def _extract_with_textrank(self) -> Dict[str, float]: - """ TextRank mode extraction """ + """TextRank mode extraction""" start_time = time.perf_counter() ranks = {} try: @@ -111,12 +117,13 @@ def _extract_with_textrank(self) -> Dict[str, float]: except MemoryError as e: log.critical("TextRank memory error (text too large?): %s", e) end_time = time.perf_counter() - log.debug("TextRank Keyword extraction time: %.2f seconds", - end_time - start_time) + log.debug( + "TextRank Keyword extraction time: %.2f seconds", end_time - start_time + ) return ranks def _extract_with_hybrid(self) -> Dict[str, float]: - """ Hybrid mode extraction """ + """Hybrid mode extraction""" ranks = {} if isinstance(llm_settings.hybrid_llm_weights, float): @@ -140,7 +147,7 @@ def _extract_with_hybrid(self) -> Dict[str, float]: if word in llm_scores: ranks[word] += llm_scores[word] * llm_weights if word in tr_scores: - ranks[word] += tr_scores[word] * (1-llm_weights) + ranks[word] += tr_scores[word] * (1 - llm_weights) end_time = time.perf_counter() log.debug("Hybrid Keyword extraction time: %.2f seconds", end_time - start_time) @@ -151,13 +158,11 @@ def _extract_keywords_from_response( response: str, lowercase: bool = True, start_token: str = "", -<<<<<<< HEAD ) -> Dict[str, float]: - results = {} # use re.escape(start_token) if start_token contains special chars like */&/^ etc. - matches = re.findall(rf'{start_token}([^\n]+\n?)', response) + matches = re.findall(rf"{start_token}([^\n]+\n?)", response) for match in matches: match = match.strip() @@ -175,34 +180,13 @@ def _extract_keywords_from_response( continue score_val = float(score_raw) if not 0.0 <= score_val <= 1.0: - log.warning("Score out of range for %s: %s", word_raw, score_val) + log.warning( + "Score out of range for %s: %s", word_raw, score_val + ) score_val = min(1.0, max(0.0, score_val)) word_out = word_raw.lower() if lowercase else word_raw results[word_out] = score_val except (ValueError, AttributeError) as e: log.warning("Failed to parse item '%s': %s", item, e) continue -======= - ) -> Set[str]: - keywords = [] - # use re.escape(start_token) if start_token contains special chars like */&/^ etc. - matches = re.findall(rf"{start_token}[^\n]+\n?", response) - - for match in matches: - match = match[len(start_token) :].strip() - keywords.extend( - k.lower() if lowercase else k - for k in re.split(r"[,，]+", match) - if len(k.strip()) > 1 - ) - - # if the keyword consists of multiple words, split into sub-words (removing stopwords) - results = set(keywords) - for token in keywords: - sub_tokens = re.findall(r"\w+", token) - if len(sub_tokens) > 1: - results.update( - w for w in sub_tokens if w not in NLTKHelper().stopwords(lang=self._language) - ) ->>>>>>> 78011d3 (Refactor: text2germlin with PCgraph framework (#50)) return results diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py similarity index 54% rename from hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py rename to hugegraph-llm/src/hugegraph_llm/operators/operator_list.py index 58848f827..6b6bf48e2 100644 --- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py +++ b/hugegraph-llm/src/hugegraph_llm/operators/operator_list.py @@ -14,45 +14,137 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from typing import Optional, List, Literal, Union - -from typing import Any, Dict, List, Literal, Optional - -from hugegraph_llm.config import huge_settings, prompt +from hugegraph_llm.config import huge_settings from hugegraph_llm.models.embeddings.base import BaseEmbedding -from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.models.llms.base import BaseLLM -from hugegraph_llm.models.llms.init_llm import LLMs -from hugegraph_llm.operators.common_op.merge_dedup_rerank import MergeDedupRerank +from hugegraph_llm.operators.common_op.check_schema import CheckSchema from hugegraph_llm.operators.common_op.print_result import PrintResult -from hugegraph_llm.operators.document_op.word_extract import WordExtract -from hugegraph_llm.operators.hugegraph_op.graph_rag_query import GraphRAGQuery from hugegraph_llm.operators.hugegraph_op.schema_manager import SchemaManager +from hugegraph_llm.operators.index_op.build_gremlin_example_index import ( + BuildGremlinExampleIndex, +) +from hugegraph_llm.operators.index_op.gremlin_example_index_query import ( + GremlinExampleIndexQuery, +) +from hugegraph_llm.operators.llm_op.gremlin_generate import GremlinGenerateSynthesize +from hugegraph_llm.utils.decorators import log_time, log_operator_time, record_rpm +from hugegraph_llm.operators.hugegraph_op.fetch_graph_data import FetchGraphData +from hugegraph_llm.operators.document_op.chunk_split import ChunkSplit +from hugegraph_llm.operators.llm_op.info_extract import InfoExtract +from hugegraph_llm.operators.llm_op.property_graph_extract import PropertyGraphExtract +from hugegraph_llm.operators.llm_op.disambiguate_data import DisambiguateData +from hugegraph_llm.operators.hugegraph_op.commit_to_hugegraph import Commit2Graph +from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex +from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex +from hugegraph_llm.operators.document_op.word_extract import WordExtract +from hugegraph_llm.operators.llm_op.keyword_extract import KeywordExtract from hugegraph_llm.operators.index_op.semantic_id_query import SemanticIdQuery from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery +from hugegraph_llm.operators.common_op.merge_dedup_rerank import MergeDedupRerank from hugegraph_llm.operators.llm_op.answer_synthesize import AnswerSynthesize -from hugegraph_llm.operators.llm_op.keyword_extract import KeywordExtract -from hugegraph_llm.utils.decorators import log_operator_time, log_time, record_rpm +from pyhugegraph.client import PyHugeClient + +class OperatorList: + def __init__( + self, + llm: BaseLLM, + embedding: BaseEmbedding, + graph: Optional[PyHugeClient] = None, + ): + self.llm = llm + self.embedding = embedding + self.result = None + self.operators = [] + self.graph = graph -class RAGPipeline: - """ - RAGPipeline is a (core)class that encapsulates a series of operations for extracting information from text, - querying graph databases and vector indices, merging and re-ranking results, and generating answers. - """ + def clear(self): + self.operators = [] + return self - def __init__(self, llm: Optional[BaseLLM] = None, embedding: Optional[BaseEmbedding] = None): - """ - Initialize the RAGPipeline with optional LLM and embedding models. + def example_index_build(self, examples): + self.operators.append(BuildGremlinExampleIndex(self.embedding, examples)) + return self - :param llm: Optional LLM model to use. - :param embedding: Optional embedding model to use. - """ - self._chat_llm = llm or LLMs().get_chat_llm() - self._extract_llm = llm or LLMs().get_extract_llm() - self._text2gqlt_llm = llm or LLMs().get_text2gql_llm() - self._embedding = embedding or Embeddings().get_embedding() - self._operators: List[Any] = [] + def import_schema( + self, from_hugegraph=None, from_extraction=None, from_user_defined=None + ): + if from_hugegraph: + self.operators.append(SchemaManager(from_hugegraph)) + elif from_user_defined: + self.operators.append(CheckSchema(from_user_defined)) + elif from_extraction: + raise NotImplementedError("Not implemented yet") + else: + raise ValueError("No input data / invalid schema type") + return self + + def example_index_query(self, num_examples): + self.operators.append(GremlinExampleIndexQuery(self.embedding, num_examples)) + return self + + def gremlin_generate_synthesize( + self, + schema, + gremlin_prompt: Optional[str] = None, + vertices: Optional[List[str]] = None, + ): + self.operators.append( + GremlinGenerateSynthesize(self.llm, schema, vertices, gremlin_prompt) + ) + return self + + def print_result(self): + self.operators.append(PrintResult()) + return self + + def fetch_graph_data(self): + if self.graph is None: + raise ValueError("graph client is required for fetch_graph_data operation") + self.operators.append(FetchGraphData(self.graph)) + return self + + def chunk_split( + self, + text: Union[str, List[str]], # text to be split + split_type: Literal["document", "paragraph", "sentence"] = "document", + language: Literal["zh", "en"] = "zh", + ): + self.operators.append(ChunkSplit(text, split_type, language)) + return self + + def extract_info( + self, + example_prompt: Optional[str] = None, + extract_type: Literal["triples", "property_graph"] = "triples", + ): + if extract_type == "triples": + self.operators.append(InfoExtract(self.llm, example_prompt)) + elif extract_type == "property_graph": + self.operators.append(PropertyGraphExtract(self.llm, example_prompt)) + else: + raise ValueError( + f"invalid extract_type: {extract_type!r}, expected 'triples' or 'property_graph'" + ) + return self + + def disambiguate_word_sense(self): + self.operators.append(DisambiguateData(self.llm)) + return self + + def commit_to_hugegraph(self): + self.operators.append(Commit2Graph()) + return self + + def build_vertex_id_semantic_index(self): + self.operators.append(BuildSemanticIndex(self.embedding)) + return self + + def build_vector_index(self): + self.operators.append(BuildVectorIndex(self.embedding)) + return self def extract_word(self, text: Optional[str] = None): """ @@ -61,7 +153,7 @@ def extract_word(self, text: Optional[str] = None): :param text: Text to extract words from. :return: Self-instance for chaining. """ - self._operators.append(WordExtract(text=text)) + self.operators.append(WordExtract(text=text)) return self def extract_keywords( @@ -76,18 +168,11 @@ def extract_keywords( :param extract_template: Template for keyword extraction. :return: Self-instance for chaining. """ - self._operators.append( - KeywordExtract( - text=text, - extract_template=extract_template - ) + self.operators.append( + KeywordExtract(text=text, extract_template=extract_template) ) return self - def import_schema(self, graph_name: str): - self._operators.append(SchemaManager(graph_name)) - return self - def keywords_to_vid( self, by: Literal["query", "keywords"] = "keywords", @@ -103,9 +188,9 @@ def keywords_to_vid( :param vector_dis_threshold: Vector distance threshold. :return: Self-instance for chaining. """ - self._operators.append( + self.operators.append( SemanticIdQuery( - embedding=self._embedding, + embedding=self.embedding, by=by, topk_per_keyword=topk_per_keyword, topk_per_query=topk_per_query, @@ -114,41 +199,6 @@ def keywords_to_vid( ) return self - def query_graphdb( - self, - max_deep: int = 2, - max_graph_items: int = huge_settings.max_graph_items, - max_v_prop_len: int = 2048, - max_e_prop_len: int = 256, - prop_to_match: Optional[str] = None, - num_gremlin_generate_example: Optional[int] = -1, - gremlin_prompt: Optional[str] = prompt.gremlin_generate_prompt, - ): - """ - Add a graph RAG query operator to the pipeline. - - :param max_deep: Maximum depth for the graph query. - :param max_graph_items: Maximum number of items to retrieve. - :param max_v_prop_len: Maximum length of vertex properties. - :param max_e_prop_len: Maximum length of edge properties. - :param prop_to_match: Property to match in the graph. - :param num_gremlin_generate_example: Number of examples to generate. - :param gremlin_prompt: Gremlin prompt for generating examples. - :return: Self-instance for chaining. - """ - self._operators.append( - GraphRAGQuery( - max_deep=max_deep, - max_graph_items=max_graph_items, - max_v_prop_len=max_v_prop_len, - max_e_prop_len=max_e_prop_len, - prop_to_match=prop_to_match, - num_gremlin_generate_example=num_gremlin_generate_example, - gremlin_prompt=gremlin_prompt, - ) - ) - return self - def query_vector_index(self, max_items: int = 3): """ Add a vector index query operator to the pipeline. @@ -156,9 +206,9 @@ def query_vector_index(self, max_items: int = 3): :param max_items: Maximum number of items to retrieve. :return: Self-instance for chaining. """ - self._operators.append( + self.operators.append( VectorIndexQuery( - embedding=self._embedding, + embedding=self.embedding, topk=max_items, ) ) @@ -177,9 +227,9 @@ def merge_dedup_rerank( :return: Self-instance for chaining. """ - self._operators.append( + self.operators.append( MergeDedupRerank( - embedding=self._embedding, + embedding=self.embedding, graph_ratio=graph_ratio, method=rerank_method, near_neighbor_first=near_neighbor_first, @@ -207,7 +257,7 @@ def synthesize_answer( :param answer_prompt: Template for the answer synthesis prompt. :return: Self-instance for chaining. """ - self._operators.append( + self.operators.append( AnswerSynthesize( raw_answer=raw_answer, vector_only_answer=vector_only_answer, @@ -218,32 +268,11 @@ def synthesize_answer( ) return self - def print_result(self): - """ - Add a print result operator to the pipeline. - - :return: Self-instance for chaining. - """ - self._operators.append(PrintResult()) - return self - @log_time("total time") @record_rpm - def run(self, **kwargs) -> Dict[str, Any]: - """ - Execute all operators in the pipeline in sequence. - - :param kwargs: Additional context to pass to operators. - :return: Final context after all operators have been executed. - """ - if len(self._operators) == 0: - self.extract_keywords().query_graphdb( - max_graph_items=kwargs.get("max_graph_items") - ).synthesize_answer() - + def run(self, **kwargs): context = kwargs - - for operator in self._operators: + for operator in self.operators: context = self._run_operator(operator, context) return context diff --git a/hugegraph-llm/src/hugegraph_llm/operators/util.py b/hugegraph-llm/src/hugegraph_llm/operators/util.py deleted file mode 100644 index 60bdc2e86..000000000 --- a/hugegraph-llm/src/hugegraph_llm/operators/util.py +++ /dev/null @@ -1,27 +0,0 @@ -# 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 PyCGraph import CStatus - - -def init_context(obj) -> CStatus: - try: - obj.context = obj.getGParamWithNoEmpty("wkflow_state") - obj.wk_input = obj.getGParamWithNoEmpty("wkflow_input") - if obj.context is None or obj.wk_input is None: - return CStatus(-1, "Required workflow parameters not found") - return CStatus() - except Exception as e: - return CStatus(-1, f"Failed to initialize context: {str(e)}") diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 3a6fd3c1c..429aba955 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py @@ -13,64 +13,71 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import AsyncGenerator, Union, List, Optional, Any, Dict from PyCGraph import GParam, CStatus -from typing import Union, List, Optional, Any +from hugegraph_llm.utils.log import log class WkFlowInput(GParam): - texts: Union[str, List[str]] = None # texts input used by ChunkSplit Node - language: str = None # language configuration used by ChunkSplit Node - split_type: str = None # split type used by ChunkSplit Node - example_prompt: str = None # need by graph information extract - schema: str = None # Schema information requeired by SchemaNode - data_json = None - extract_type = None - query_examples = None - few_shot_schema = None + texts: Optional[Union[str, List[str]]] = None # texts input used by ChunkSplit Node + 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 + schema: Optional[str] = None # Schema information requeired by SchemaNode + data_json: Optional[Dict[str, Any]] = None + extract_type: Optional[str] = None + query_examples: Optional[Any] = None + few_shot_schema: Optional[Any] = None # Fields related to PromptGenerate - source_text: str = None # Original text - scenario: str = None # Scenario description - example_name: str = None # Example name + source_text: Optional[str] = None # Original text + scenario: Optional[str] = None # Scenario description + example_name: Optional[str] = None # Example name # Fields for Text2Gremlin - example_num: int = None - gremlin_prompt: str = None + example_num: Optional[int] = None requested_outputs: Optional[List[str]] = None # RAG Flow related fields - query: str = None # User query for RAG - vector_search: bool = None # Enable vector search - graph_search: bool = None # Enable graph search - raw_answer: bool = None # Return raw answer - vector_only_answer: bool = None # Vector only answer mode - graph_only_answer: bool = None # Graph only answer mode - graph_vector_answer: bool = None # Combined graph and vector answer - graph_ratio: float = None # Graph ratio for merging - rerank_method: str = None # Reranking method - near_neighbor_first: bool = None # Near neighbor first flag - custom_related_information: str = None # Custom related information - answer_prompt: str = None # Answer generation prompt - keywords_extract_prompt: str = None # Keywords extraction prompt - gremlin_tmpl_num: int = None # Gremlin template number - gremlin_prompt: str = None # Gremlin generation prompt - max_graph_items: int = None # Maximum graph items - topk_return_results: int = None # Top-k return results - vector_dis_threshold: float = None # Vector distance threshold - topk_per_keyword: int = None # Top-k per keyword - max_keywords: int = None - max_items: int = None + query: Optional[str] = None # User query for RAG + vector_search: Optional[bool] = None # Enable vector search + graph_search: Optional[bool] = None # Enable graph search + raw_answer: Optional[bool] = None # Return raw answer + vector_only_answer: Optional[bool] = None # Vector only answer mode + graph_only_answer: Optional[bool] = None # Graph only answer mode + graph_vector_answer: Optional[bool] = None # Combined graph and vector answer + graph_ratio: Optional[float] = None # Graph ratio for merging + rerank_method: Optional[str] = None # Reranking method + near_neighbor_first: Optional[bool] = None # Near neighbor first flag + custom_related_information: Optional[str] = None # Custom related information + answer_prompt: Optional[str] = None # Answer generation prompt + keywords_extract_prompt: Optional[str] = None # Keywords extraction prompt + gremlin_tmpl_num: Optional[int] = None # Gremlin template number + gremlin_prompt: Optional[str] = None # Gremlin generation prompt + max_graph_items: Optional[int] = None # Maximum graph items + topk_return_results: Optional[int] = None # Top-k return results + vector_dis_threshold: Optional[float] = None # Vector distance threshold + topk_per_keyword: Optional[int] = None # Top-k per keyword + max_keywords: Optional[int] = None + max_items: Optional[int] = None # Semantic query related fields - semantic_by: str = None # Semantic query method - topk_per_query: int = None # Top-k per query + semantic_by: Optional[str] = None # Semantic query method + topk_per_query: Optional[int] = None # Top-k per query # Graph query related fields - max_deep: int = None # Maximum depth for graph traversal - max_v_prop_len: int = None # Maximum vertex property length - max_e_prop_len: int = None # Maximum edge property length - prop_to_match: str = None # Property to match + max_deep: Optional[int] = None # Maximum depth for graph traversal + max_v_prop_len: Optional[int] = None # Maximum vertex property length + max_e_prop_len: Optional[int] = None # Maximum edge property length + prop_to_match: Optional[str] = None # Property to match - stream: bool = None # used for recognize stream mode + stream: Optional[bool] = None # used for recognize stream mode + + # used for rag_recall api + is_graph_rag_recall: bool = False + is_vector_only: bool = False + + # used for build text2gremin index + examples: Optional[List[Dict[str, str]]] = None def reset(self, _: CStatus) -> None: self.texts = None @@ -78,7 +85,6 @@ def reset(self, _: CStatus) -> None: self.split_type = None self.example_prompt = None self.schema = None - self.graph_name = None self.data_json = None self.extract_type = None self.query_examples = None @@ -106,7 +112,6 @@ def reset(self, _: CStatus) -> None: self.answer_prompt = None self.keywords_extract_prompt = None self.gremlin_tmpl_num = None - self.gremlin_prompt = None self.max_graph_items = None self.topk_return_results = None self.vector_dis_threshold = None @@ -123,6 +128,10 @@ def reset(self, _: CStatus) -> None: self.prop_to_match = None self.stream = None + self.examples = None + self.is_graph_rag_recall = False + self.is_vector_only = False + class WkFlowState(GParam): schema: Optional[str] = None # schema message @@ -134,9 +143,9 @@ class WkFlowState(GParam): call_count: Optional[int] = None keywords: Optional[List[str]] = None - vector_result = None - graph_result = None - keywords_embeddings = None + vector_result: Optional[Any] = None + graph_result: Optional[Any] = None + keywords_embeddings: Optional[Any] = None generated_extract_prompt: Optional[str] = None # Fields for Text2Gremlin results @@ -146,18 +155,43 @@ class WkFlowState(GParam): template_exec_res: Optional[Any] = None raw_exec_res: Optional[Any] = None - match_vids = None - vector_result = None - graph_result = None + match_vids: Optional[Any] = None + + raw_answer: Optional[str] = None + vector_only_answer: Optional[str] = None + graph_only_answer: Optional[str] = None + graph_vector_answer: Optional[str] = None + + merged_result: Optional[Any] = None + + vertex_num: Optional[int] = None + edge_num: Optional[int] = None + note: Optional[str] = None + removed_vid_vector_num: Optional[int] = None + added_vid_vector_num: Optional[int] = None + raw_texts: Optional[List] = None + query_examples: Optional[List] = None + few_shot_schema: Optional[Dict] = None + source_text: Optional[str] = None + scenario: Optional[str] = None + example_name: Optional[str] = None - raw_answer: str = None - vector_only_answer: str = None - graph_only_answer: str = None - graph_vector_answer: str = None + graph_ratio: Optional[float] = None + query: Optional[str] = None + vector_search: Optional[bool] = None + graph_search: Optional[bool] = None + max_graph_items: Optional[int] = None + stream_generator: Optional[AsyncGenerator] = None - merged_result = None + graph_result_flag: Optional[int] = None + vertex_degree_list: Optional[List] = None + knowledge_with_degree: Optional[Dict] = None + graph_context_head: Optional[str] = None - def setup(self): + embed_dim: Optional[int] = None + is_graph_rag_recall: Optional[bool] = None + + def setup(self) -> CStatus: self.schema = None self.simple_schema = None self.chunks = None @@ -184,9 +218,36 @@ def setup(self): self.graph_only_answer = None self.graph_vector_answer = None - self.vector_result = None - self.graph_result = None self.merged_result = None + + self.match_vids = None + self.vertex_num = None + self.edge_num = None + self.note = None + self.removed_vid_vector_num = None + self.added_vid_vector_num = None + + self.raw_texts = None + self.query_examples = None + self.few_shot_schema = None + self.source_text = None + self.scenario = None + self.example_name = None + + self.graph_ratio = None + self.query = None + self.vector_search = None + self.graph_search = None + self.max_graph_items = None + + self.stream_generator = None + self.graph_result_flag = None + self.vertex_degree_list = None + self.knowledge_with_degree = None + self.graph_context_head = None + + self.embed_dim = None + self.is_graph_rag_recall = None return CStatus() def to_json(self): @@ -210,4 +271,9 @@ def assign_from_json(self, data_json: dict): Assigns each key in the input json object as a member variable of WkFlowState. """ for k, v in data_json.items(): - setattr(self, k, v) + if hasattr(self, k): + setattr(self, k, v) + else: + log.warning( + "key %s should be a member of WkFlowState & type %s", k, type(v) + ) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 3f527f2fa..9c53e8183 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -16,29 +16,28 @@ # under the License. -import json import os import traceback -from typing import Dict, Any, Union, Optional +from typing import Dict, Any, Union, List import gradio as gr +from hugegraph_llm.flows import FlowName from hugegraph_llm.flows.scheduler import SchedulerSingleton +from pyhugegraph.client import PyHugeClient from .embedding_utils import get_filename_prefix, get_index_folder_name -from .hugegraph_utils import get_hg_client, clean_hg_data +from .hugegraph_utils import clean_hg_data from .log import log from .vector_index_utils import read_documents from ..config import resource_path, huge_settings, llm_settings from ..indices.vector_index import VectorIndex from ..models.embeddings.init_embedding import Embeddings -from ..models.llms.init_llm import LLMs -from ..operators.kg_construction_task import KgBuilder def get_graph_index_info(): try: scheduler = SchedulerSingleton.get_instance() - return scheduler.schedule_flow("get_graph_index_info") + return scheduler.schedule_flow(FlowName.GET_GRAPH_INDEX_INFO) except Exception as e: # pylint: disable=broad-exception-caught log.error(e) raise gr.Error(str(e)) @@ -63,63 +62,31 @@ def clean_all_graph_index(): gr.Info("Clear graph index and text2gql index successfully!") -def clean_all_graph_data(): - clean_hg_data() - log.warning("Clear graph data successfully!") - gr.Info("Clear graph data successfully!") - - -def parse_schema(schema: str, builder: KgBuilder) -> Optional[str]: - schema = schema.strip() - if schema.startswith("{"): - try: - schema = json.loads(schema) - builder.import_schema(from_user_defined=schema) - except json.JSONDecodeError: - log.error("Invalid JSON format in schema. Please check it again.") - return "ERROR: Invalid JSON format in schema. Please check it carefully." +def get_vertex_details( + vertex_ids: List[str], context: Dict[str, Any] +) -> List[Dict[str, Any]]: + if isinstance(context.get("graph_client"), PyHugeClient): + client = context["graph_client"] else: - log.info("Get schema '%s' from graphdb.", schema) - builder.import_schema(from_hugegraph=schema) - return None + url = context.get("url") or "http://localhost:8080" + graph = context.get("graph") or "hugegraph" + user = context.get("user") or "admin" + pwd = context.get("pwd") or "admin" + gs = context.get("graphspace") or None + client = PyHugeClient(url, graph, user, pwd, gs) + if not vertex_ids: + return [] + formatted_ids = ", ".join(f"'{vid}'" for vid in vertex_ids) + gremlin_query = f"g.V({formatted_ids}).limit(20)" + result = client.gremlin().exec(gremlin=gremlin_query)["data"] + return result -def extract_graph_origin(input_file, input_text, schema, example_prompt) -> str: - texts = read_documents(input_file, input_text) - builder = KgBuilder( - LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client() - ) - if not schema: - return "ERROR: please input with correct schema/format." - - error_message = parse_schema(schema, builder) - if error_message: - return error_message - builder.chunk_split(texts, "document", "zh").extract_info( - example_prompt, "property_graph" - ) - try: - context = builder.run() - if not context["vertices"] and not context["edges"]: - log.info("Please check the schema.(The schema may not match the Doc)") - return json.dumps( - { - "vertices": context["vertices"], - "edges": context["edges"], - "warning": "The schema may not match the Doc", - }, - ensure_ascii=False, - indent=2, - ) - return json.dumps( - {"vertices": context["vertices"], "edges": context["edges"]}, - ensure_ascii=False, - indent=2, - ) - except Exception as e: # pylint: disable=broad-exception-caught - log.error(e) - raise gr.Error(str(e)) +def clean_all_graph_data(): + clean_hg_data() + log.warning("Clear graph data successfully!") + gr.Info("Clear graph data successfully!") def extract_graph(input_file, input_text, schema, example_prompt) -> str: @@ -130,7 +97,7 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: try: return scheduler.schedule_flow( - "graph_extract", schema, texts, example_prompt, "property_graph" + FlowName.GRAPH_EXTRACT, schema, texts, example_prompt, "property_graph" ) except Exception as e: # pylint: disable=broad-exception-caught log.error(e) @@ -140,7 +107,7 @@ def extract_graph(input_file, input_text, schema, example_prompt) -> str: def update_vid_embedding(): scheduler = SchedulerSingleton.get_instance() try: - return scheduler.schedule_flow("update_vid_embeddings") + return scheduler.schedule_flow(FlowName.UPDATE_VID_EMBEDDINGS) except Exception as e: # pylint: disable=broad-exception-caught log.error(e) raise gr.Error(str(e)) @@ -149,7 +116,7 @@ def update_vid_embedding(): def import_graph_data(data: str, schema: str) -> Union[str, Dict[str, Any]]: try: scheduler = SchedulerSingleton.get_instance() - return scheduler.schedule_flow("import_graph_data", data, schema) + return scheduler.schedule_flow(FlowName.IMPORT_GRAPH_DATA, data, schema) except Exception as e: # pylint: disable=W0718 log.error(e) traceback.print_exc() @@ -162,7 +129,8 @@ def build_schema(input_text, query_example, few_shot): scheduler = SchedulerSingleton.get_instance() try: return scheduler.schedule_flow( - "build_schema", input_text, query_example, few_shot + FlowName.BUILD_SCHEMA, input_text, query_example, few_shot ) - except (TypeError, ValueError) as e: + except Exception as e: # pylint: disable=broad-exception-caught + log.error("Schema generation failed: %s", e) raise gr.Error(f"Schema generation failed: {e}") diff --git a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py index 301a6bdab..67904a445 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -22,6 +22,7 @@ import gradio as gr from hugegraph_llm.config import resource_path, huge_settings, llm_settings +from hugegraph_llm.flows import FlowName from hugegraph_llm.indices.vector_index import VectorIndex from hugegraph_llm.models.embeddings.init_embedding import model_map from hugegraph_llm.flows.scheduler import SchedulerSingleton @@ -50,7 +51,9 @@ def read_documents(input_file, input_text): texts.append(text) elif full_path.endswith(".pdf"): # TODO: support PDF file - raise gr.Error("PDF will be supported later! Try to upload text/docx now") + raise gr.Error( + "PDF will be supported later! Try to upload text/docx now" + ) else: raise gr.Error("Please input txt or docx file.") else: @@ -60,7 +63,9 @@ def read_documents(input_file, input_text): # pylint: disable=C0301 def get_vector_index_info(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) filename_prefix = get_filename_prefix( llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) ) @@ -87,11 +92,15 @@ def get_vector_index_info(): def clean_vector_index(): - folder_name = get_index_folder_name(huge_settings.graph_name, huge_settings.graph_space) + folder_name = get_index_folder_name( + huge_settings.graph_name, huge_settings.graph_space + ) filename_prefix = get_filename_prefix( llm_settings.embedding_type, model_map.get(llm_settings.embedding_type) ) - VectorIndex.clean(str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix) + VectorIndex.clean( + str(os.path.join(resource_path, folder_name, "chunks")), filename_prefix + ) gr.Info("Clean vector index successfully!") @@ -100,4 +109,4 @@ def build_vector_index(input_file, input_text): raise gr.Error("Please only choose one between file and text.") texts = read_documents(input_file, input_text) scheduler = SchedulerSingleton.get_instance() - return scheduler.schedule_flow("build_vector_index", texts) + return scheduler.schedule_flow(FlowName.BUILD_VECTOR_INDEX, texts) diff --git a/hugegraph-ml/pyproject.toml b/hugegraph-ml/pyproject.toml index 6d46ba74c..929eb3aa1 100644 --- a/hugegraph-ml/pyproject.toml +++ b/hugegraph-ml/pyproject.toml @@ -22,7 +22,7 @@ build-backend = "hatchling.build" [project] name = "hugegraph-ml" -version = "1.5.0" +version = "1.7.0" description = "Machine learning extensions for Apache HugeGraph." authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, diff --git a/hugegraph-python-client/pyproject.toml b/hugegraph-python-client/pyproject.toml index 81565d9ab..ddae125d8 100644 --- a/hugegraph-python-client/pyproject.toml +++ b/hugegraph-python-client/pyproject.toml @@ -17,7 +17,7 @@ [project] name = "hugegraph-python-client" -version = "1.5.0" +version = "1.7.0" description = "A Python SDK for Apache HugeGraph Database." authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, diff --git a/pyproject.toml b/pyproject.toml index 8bcf58929..2dd4161d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ [project] name = "hugegraph-ai" -version = "1.5.0" +version = "1.7.0" description = "A repository for AI-related projects for Apache HugeGraph." authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" }, diff --git a/scripts/build_llm_image.sh b/scripts/build_llm_image.sh old mode 100644 new mode 100755 index 42aa36e39..7425b3df9 --- a/scripts/build_llm_image.sh +++ b/scripts/build_llm_image.sh @@ -18,7 +18,7 @@ set -e -tag="1.5.0" +tag="1.7.0" script_dir=$(realpath "$(dirname "$0")") diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/design.md b/spec/hugegraph-llm/fixed_flow/design.md similarity index 99% rename from .vibedev/spec/hugegraph-llm/fixed_flow/design.md rename to spec/hugegraph-llm/fixed_flow/design.md index c5777236d..5ad64407a 100644 --- a/.vibedev/spec/hugegraph-llm/fixed_flow/design.md +++ b/spec/hugegraph-llm/fixed_flow/design.md @@ -202,7 +202,7 @@ flowchart TD - `BuildVectorIndexFlow`: 向量索引构建工作流 - `GraphExtractFlow`: 图抽取工作流 - `ImportGraphDataFlow`: 图数据导入工作流 - - `UpdateVidEmbeddingsFlows`: 向量更新工作流 + - `UpdateVidEmbeddingsFlow`: 向量更新工作流 - `GetGraphIndexInfoFlow`: 图索引信息获取工作流 - `BuildSchemaFlow`: 模式构建工作流 - `PromptGenerateFlow`: 提示词生成工作流 @@ -407,7 +407,6 @@ class GraphExtractFlow(BaseFlow): prepared_input.split_type = "document" prepared_input.example_prompt = example_prompt prepared_input.schema = schema - prepare_schema(prepared_input, schema) return def build_flow(self, schema, texts, example_prompt, extract_type): diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/requirements.md b/spec/hugegraph-llm/fixed_flow/requirements.md similarity index 100% rename from .vibedev/spec/hugegraph-llm/fixed_flow/requirements.md rename to spec/hugegraph-llm/fixed_flow/requirements.md diff --git a/.vibedev/spec/hugegraph-llm/fixed_flow/tasks.md b/spec/hugegraph-llm/fixed_flow/tasks.md similarity index 100% rename from .vibedev/spec/hugegraph-llm/fixed_flow/tasks.md rename to spec/hugegraph-llm/fixed_flow/tasks.md diff --git a/style/pylint.conf b/style/pylint.conf index 6ccb7a078..4fb3a17c2 100644 --- a/style/pylint.conf +++ b/style/pylint.conf @@ -476,6 +476,7 @@ disable=raw-checker-failed, # it should appear only once). See also the "--disable" option for examples. enable= +extension-pkg-whitelist=PyCGraph [METHOD_ARGS] @@ -596,7 +597,8 @@ contextmanager-decorators=contextlib.contextmanager # List of members which are set dynamically and missed by pylint inference # system, and so shouldn't trigger E1101 when accessed. Python regular # expressions are accepted. -generated-members= +ignored-modules=PyCGraph +generated-members=PyCGraph.* # Tells whether to warn about missing members when the owner of the attribute # is inferred to be None. diff --git a/vermeer-python-client/pyproject.toml b/vermeer-python-client/pyproject.toml index 986010899..d60acc075 100644 --- a/vermeer-python-client/pyproject.toml +++ b/vermeer-python-client/pyproject.toml @@ -17,7 +17,7 @@ [project] name = "vermeer-python-client" -version = "1.5.0" # Independently managed version for the vermeer-python-client package +version = "1.7.0" # Independently managed version for the vermeer-python-client package description = "A Python client library for interacting with Vermeer, a tool for managing and analyzing large-scale graph data." authors = [ { name = "Apache HugeGraph Contributors", email = "dev@hugegraph.apache.org" } @@ -33,7 +33,7 @@ dependencies = [ "setuptools", "urllib3", "rich", - + # Vermeer specific dependencies "python-dateutil", ] From 8bbd29b489e9e60893c977eef3edd14f1fb4dd9b Mon Sep 17 00:00:00 2001 From: lingxiao Date: Mon, 27 Oct 2025 16:10:56 +0800 Subject: [PATCH 61/71] chore: commit staged changes before merge --- hugegraph-python-client/src/pyhugegraph/utils/util.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/hugegraph-python-client/src/pyhugegraph/utils/util.py b/hugegraph-python-client/src/pyhugegraph/utils/util.py index 76770a818..1f18b8f19 100644 --- a/hugegraph-python-client/src/pyhugegraph/utils/util.py +++ b/hugegraph-python-client/src/pyhugegraph/utils/util.py @@ -58,7 +58,12 @@ def check_if_success(response, error=None): req = response.request req_body = req.body if req.body else "Empty body" response_body = response.text if response.text else "Empty body" - log.error() + log.error( + "Error-Client: Request URL: %s, Request Body: %s, Response Body: %s", + getattr(req, "url", "Unknown URL"), + req_body, + response_body, + ) raise error return True From 4ba51ca4539b1d5314ed9d05839016004f112125 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Mon, 27 Oct 2025 16:43:45 +0800 Subject: [PATCH 62/71] fix: update GetGraphIndexInfoFlow to use new vector index API - Replace VectorIndex.from_index_file() with VectorStoreBase.from_name() - Use get_vector_index_class() factory function - Use get_vector_index_info() method for consistent interface - Fix ImportError after merge --- .../flows/get_graph_index_info.py | 40 ++++++++----------- 1 file changed, 16 insertions(+), 24 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py index 86d08bf2d..5c09eaef0 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py @@ -14,20 +14,14 @@ # limitations under the License. import json -import os from PyCGraph import GPipeline -from hugegraph_llm.config import huge_settings, llm_settings, resource_path +from hugegraph_llm.config import huge_settings, index_settings from hugegraph_llm.flows.common import BaseFlow -from hugegraph_llm.indices.vector_index import VectorIndex -from hugegraph_llm.models.embeddings.init_embedding import model_map +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState from hugegraph_llm.nodes.hugegraph_node.fetch_graph_data import FetchGraphDataNode -from hugegraph_llm.utils.embedding_utils import ( - get_filename_prefix, - get_index_folder_name, -) # pylint: disable=arguments-differ,keyword-arg-before-vararg @@ -49,22 +43,20 @@ def build_flow(self, **kwargs): return pipeline def post_deal(self, pipeline=None): + # Lazy import to avoid circular dependency + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel + graph_summary_info = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - index_dir = str(os.path.join(resource_path, folder_name, "graph_vids")) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, - model_map.get(llm_settings.embedding_type, None), - ) + try: - vector_index = VectorIndex.from_index_file(index_dir, filename_prefix) - except (RuntimeError, OSError): - return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) - graph_summary_info["vid_index"] = { - "embed_dim": vector_index.index.d, - "num_vectors": vector_index.index.ntotal, - "num_vids": len(vector_index.properties), - } + vector_index_class = get_vector_index_class(index_settings.cur_vector_index) + embedding = Embeddings().get_embedding() + vector_index = vector_index_class.from_name( + embedding.get_embedding_dim(), huge_settings.graph_name, "graph_vids" + ) + graph_summary_info["vid_index"] = vector_index.get_vector_index_info() + except Exception: # pylint: disable=broad-except + # If vector index doesn't exist or fails to load, just return graph summary + pass + return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) From 66b096171ea65558f83bd6cab2f6a7c2a25f53f9 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Mon, 27 Oct 2025 16:44:50 +0800 Subject: [PATCH 63/71] fix: update graph_index_utils to use new vector index API - Remove VectorIndex import and usage - Update clean_all_graph_index() to use get_vector_index_class() - Use VectorStoreBase.clean() method with graph_name parameter - Fix remaining ImportError after merge --- .../hugegraph_llm/utils/graph_index_utils.py | 26 ++++++------------- 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 9c53e8183..b9b29c384 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -25,13 +25,10 @@ from hugegraph_llm.flows.scheduler import SchedulerSingleton from pyhugegraph.client import PyHugeClient -from .embedding_utils import get_filename_prefix, get_index_folder_name from .hugegraph_utils import clean_hg_data from .log import log from .vector_index_utils import read_documents -from ..config import resource_path, huge_settings, llm_settings -from ..indices.vector_index import VectorIndex -from ..models.embeddings.init_embedding import Embeddings +from ..config import huge_settings def get_graph_index_info(): @@ -44,20 +41,13 @@ def get_graph_index_info(): def clean_all_graph_index(): - folder_name = get_index_folder_name( - huge_settings.graph_name, huge_settings.graph_space - ) - filename_prefix = get_filename_prefix( - llm_settings.embedding_type, - getattr(Embeddings().get_embedding(), "model_name", None), - ) - VectorIndex.clean( - str(os.path.join(resource_path, folder_name, "graph_vids")), filename_prefix - ) - VectorIndex.clean( - str(os.path.join(resource_path, folder_name, "gremlin_examples")), - filename_prefix, - ) + # Lazy import to avoid circular dependency + from .vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel + from ..config import index_settings # pylint: disable=import-outside-toplevel + + vector_index = get_vector_index_class(index_settings.cur_vector_index) + vector_index.clean(huge_settings.graph_name, "graph_vids") + vector_index.clean("gremlin_examples") log.warning("Clear graph index and text2gql index successfully!") gr.Info("Clear graph index and text2gql index successfully!") From 9690da7589f0408bd3d7e0be076e2b43f1b7f51d Mon Sep 17 00:00:00 2001 From: lingxiao Date: Mon, 27 Oct 2025 17:42:31 +0800 Subject: [PATCH 64/71] fix conflicts --- .../demo/rag_demo/configs_block.py | 127 +++++++++++++++++- .../hugegraph_llm/models/embeddings/base.py | 8 +- .../models/embeddings/init_embedding.py | 48 +++++-- .../models/embeddings/litellm.py | 74 +++++++--- .../hugegraph_llm/models/embeddings/ollama.py | 19 ++- .../hugegraph_llm/models/embeddings/openai.py | 33 +++-- .../index_node/build_gremlin_example_index.py | 4 +- .../nodes/index_node/build_semantic_index.py | 4 +- .../nodes/index_node/build_vector_index.py | 4 +- .../index_node/gremlin_example_index_query.py | 4 +- .../index_node/semantic_id_query_node.py | 6 +- .../nodes/index_node/vector_query_node.py | 4 +- 12 files changed, 270 insertions(+), 65 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 8c595c30d..4bf875b9c 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -25,7 +25,7 @@ from dotenv import dotenv_values from requests.auth import HTTPBasicAuth -from hugegraph_llm.config import huge_settings, llm_settings +from hugegraph_llm.config import huge_settings, index_settings, llm_settings from hugegraph_llm.models.embeddings.litellm import LiteLLMEmbedding from hugegraph_llm.models.llms.litellm import LiteLLMClient from hugegraph_llm.utils.log import log @@ -106,6 +106,83 @@ def test_api_connection( return resp.status_code +def apply_vector_engine(engine: str): + # Persist the vector engine selection + setattr(index_settings, "cur_vector_index", engine) + try: + index_settings.update_env() + except Exception: # pylint: disable=W0718 + pass + gr.Info("Configured!") + + +def apply_vector_engine_backend( + engine: str, + host: Optional[str] = None, + port: Optional[str] = None, + user: Optional[str] = None, + password: Optional[str] = None, + api_key: Optional[str] = None, + origin_call=None, +) -> int: + """Test connection and persist per-engine connection settings""" + status_code = -1 + + # Test connection first + try: + if engine == "Milvus": + from pymilvus import connections, utility + + connections.connect( + host=host, port=int(port or 19530), user=user or "", password=password or "" + ) + # Test if we can list collections + _ = utility.list_collections() + connections.disconnect("default") + status_code = 200 + elif engine == "Qdrant": + from qdrant_client import QdrantClient + + client = QdrantClient(host=host, port=int(port or 6333), api_key=api_key) + # Test if we can get collections + _ = client.get_collections() + status_code = 200 + except ImportError as e: + msg = f"Missing dependency: {e}. Please install with: uv sync --extra vectordb" + if origin_call is None: + raise gr.Error(msg) from e + return -1 + except Exception as e: + msg = f"Connection failed: {e}" + log.error(msg) + if origin_call is None: + raise gr.Error(msg) from e + return -1 + + # Persist settings after successful test + if engine == "Milvus": + if host is not None: + index_settings.milvus_host = host + if port is not None and str(port).strip(): + index_settings.milvus_port = int(port) # type: ignore[arg-type] + index_settings.milvus_user = user or "" + index_settings.milvus_password = password or "" + elif engine == "Qdrant": + if host is not None: + index_settings.qdrant_host = host + if port is not None and str(port).strip(): + index_settings.qdrant_port = int(port) # type: ignore[arg-type] + # Empty string treated as None for api key + index_settings.qdrant_api_key = api_key or None + + try: + index_settings.update_env() + except Exception: # pylint: disable=W0718 + pass + gr.Info("Configured!") + return status_code + + def apply_embedding_config(arg1, arg2, arg3, origin_call=None) -> int: status_code = -1 embedding_option = llm_settings.embedding_type @@ -653,6 +730,54 @@ def reranker_settings(reranker_type): inputs=reranker_config_input, # pylint: disable=no-member ) + with gr.Accordion("5. Set up the vector engine.", open=False): + engine_selector = gr.Dropdown( + choices=["Faiss", "Milvus", "Qdrant"], + value=index_settings.cur_vector_index, + label="Select vector engine.", + ) + engine_selector.select( + fn=lambda engine: setattr(index_settings, "cur_vector_index", engine), + inputs=[engine_selector], + ) + + @gr.render(inputs=[engine_selector]) + def vector_engine_settings(engine): + if engine == "Milvus": + with gr.Row(): + milvus_inputs = [ + gr.Textbox(value=index_settings.milvus_host, label="host"), + gr.Textbox(value=str(index_settings.milvus_port), label="port"), + gr.Textbox(value=index_settings.milvus_user, label="user"), + gr.Textbox( + value=index_settings.milvus_password, label="password", type="password" + ), + ] + apply_backend_button = gr.Button("Apply Configuration") + apply_backend_button.click( + partial(apply_vector_engine_backend, "Milvus"), inputs=milvus_inputs + ) + elif engine == "Qdrant": + with gr.Row(): + qdrant_inputs = [ + gr.Textbox(value=index_settings.qdrant_host, label="host"), + gr.Textbox(value=str(index_settings.qdrant_port), label="port"), + gr.Textbox( + value=(index_settings.qdrant_api_key or ""), + label="api_key", + type="password", + ), + ] + apply_backend_button = gr.Button("Apply Configuration") + apply_backend_button.click( + lambda h, p, k: apply_vector_engine_backend("Qdrant", h, p, None, None, k), + inputs=qdrant_inputs, + ) + else: + gr.Markdown("✅ Faiss 本地索引无需额外配置。") + apply_faiss_button = gr.Button("Apply Configuration") + apply_faiss_button.click(lambda: apply_vector_engine(engine)) + # The reason for returning this partial value is the functional need to refresh the ui return graph_config_input diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py index a96650725..e25d61c4f 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py @@ -67,8 +67,8 @@ def get_embedding_dim( """Get the dimension of the embedding.""" @abstractmethod - def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts in a single batch. + def get_texts_embeddings(self, texts: List[str], batch_size: int = 32) -> List[List[float]]: + """Get embeddings for multiple texts with automatic batch splitting. This method should efficiently process multiple texts at once by leveraging the embedding model's batching capabilities, which is typically more efficient @@ -87,8 +87,8 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: """ @abstractmethod - async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts in a single batch asynchronously. + async def async_get_texts_embeddings(self, texts: List[str], batch_size: int = 32) -> List[List[float]]: + """Get embeddings for multiple texts asynchronously with automatic batch splitting. This method should efficiently process multiple texts at once by leveraging the embedding model's batching capabilities, which is typically more efficient diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index 19d9d6084..bd903f049 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -57,28 +57,58 @@ def __init__(self): self.embedding_type = llm_settings.embedding_type def get_embedding(self): + """Get embedding instance and dynamically determine dimension if needed.""" if self.embedding_type == "openai": - assert llm_settings.openai_embedding_model_dim, "openai_embedding_model_dim is need" - return OpenAIEmbedding( - embedding_dimension=llm_settings.openai_embedding_model_dim, + # Create with default dimension first + embedding = OpenAIEmbedding( model_name=llm_settings.openai_embedding_model, api_key=llm_settings.openai_embedding_api_key, api_base=llm_settings.openai_embedding_api_base, ) + # Dynamically get actual dimension + try: + test_vec = embedding.get_text_embedding("test") + embedding.embedding_dimension = len(test_vec) + except Exception: # pylint: disable=broad-except + pass # Keep default dimension + return embedding if self.embedding_type == "ollama/local": - assert llm_settings.ollama_embedding_model_dim, "ollama_embedding_model_dim is need" - return OllamaEmbedding( - embedding_dimension=llm_settings.ollama_embedding_model_dim, + # Create with default dimension first + embedding = OllamaEmbedding( model=llm_settings.ollama_embedding_model, host=llm_settings.ollama_embedding_host, port=llm_settings.ollama_embedding_port, ) + # Dynamically get actual dimension + try: + test_vec = embedding.get_text_embedding("test") + embedding.embedding_dimension = len(test_vec) + except Exception: # pylint: disable=broad-except + pass # Keep default dimension + return embedding if self.embedding_type == "litellm": - return LiteLLMEmbedding( - embedding_dimension=llm_settings.litellm_embedding_model_dim, + # For LiteLLM, we need to get dimension dynamically + # Create a temporary instance to test dimension + temp_embedding = LiteLLMEmbedding( + embedding_dimension=1536, # Temporary default model_name=llm_settings.litellm_embedding_model, api_key=llm_settings.litellm_embedding_api_key, api_base=llm_settings.litellm_embedding_api_base, - ) # type: ignore + ) + # Get actual dimension + try: + test_vec = temp_embedding.get_text_embedding("test") + actual_dim = len(test_vec) + except Exception: # pylint: disable=broad-except + actual_dim = 1536 # Fallback + + # Create final instance with correct dimension + embedding = LiteLLMEmbedding( + embedding_dimension=actual_dim, + model_name=llm_settings.litellm_embedding_model, + api_key=llm_settings.litellm_embedding_api_key, + api_base=llm_settings.litellm_embedding_api_base, + ) + return embedding # type: ignore raise Exception("embedding type is not supported !") diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py index c5effaacf..3f9619cd9 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/litellm.py @@ -64,17 +64,34 @@ def get_text_embedding(self, text: str) -> List[float]: log.error("Error in LiteLLM embedding call: %s", e) raise - def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts.""" + def get_texts_embeddings(self, texts: List[str], batch_size: int = 32) -> List[List[float]]: + """Get embeddings for multiple texts with automatic batch splitting. + + Parameters + ---------- + texts : List[str] + A list of text strings to be embedded. + batch_size : int, optional + Maximum number of texts to process in a single API call (default: 32). + + Returns + ------- + List[List[float]] + A list of embedding vectors. + """ + all_embeddings = [] try: - response = embedding( - model=self.model, - input=texts, - api_key=self.api_key, - api_base=self.api_base, - ) - log.info("Token usage: %s", response.usage) - return [data["embedding"] for data in response.data] + for i in range(0, len(texts), batch_size): + batch = texts[i:i + batch_size] + response = embedding( + model=self.model, + input=batch, + api_key=self.api_key, + api_base=self.api_base, + ) + log.info("Token usage: %s", response.usage) + all_embeddings.extend([data["embedding"] for data in response.data]) + return all_embeddings except (RateLimitError, APIConnectionError, APIError) as e: log.error("Error in LiteLLM batch embedding call: %s", e) raise @@ -94,17 +111,34 @@ async def async_get_text_embedding(self, text: str) -> List[float]: log.error("Error in async LiteLLM embedding call: %s", e) raise - async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts asynchronously.""" + async def async_get_texts_embeddings(self, texts: List[str], batch_size: int = 32) -> List[List[float]]: + """Get embeddings for multiple texts asynchronously with automatic batch splitting. + + Parameters + ---------- + texts : List[str] + A list of text strings to be embedded. + batch_size : int, optional + Maximum number of texts to process in a single API call (default: 32). + + Returns + ------- + List[List[float]] + A list of embedding vectors. + """ + all_embeddings = [] try: - response = await aembedding( - model=self.model, - input=texts, - api_key=self.api_key, - api_base=self.api_base, - ) - log.info("Token usage: %s", response.usage) - return [data["embedding"] for data in response.data] + for i in range(0, len(texts), batch_size): + batch = texts[i:i + batch_size] + response = await aembedding( + model=self.model, + input=batch, + api_key=self.api_key, + api_base=self.api_base, + ) + log.info("Token usage: %s", response.usage) + all_embeddings.extend([data["embedding"] for data in response.data]) + return all_embeddings except (RateLimitError, APIConnectionError, APIError) as e: log.error("Error in async LiteLLM embedding call: %s", e) raise diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py index c02590695..910372b43 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py @@ -45,17 +45,18 @@ def get_text_embedding(self, text: str) -> List[float]: """Comment""" return list(self.client.embed(model=self.model, input=text)["embeddings"][0]) - def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts in a single batch. + def get_texts_embeddings(self, texts: List[str], batch_size: int = 32) -> List[List[float]]: + """Get embeddings for multiple texts with automatic batch splitting. - This method efficiently processes multiple texts at once by leveraging - Ollama's batching capabilities, which is more efficient than processing - texts individually. + This method efficiently processes multiple texts by splitting them into + smaller batches to respect API rate limits and batch size constraints. Parameters ---------- texts : List[str] A list of text strings to be embedded. + batch_size : int, optional + Maximum number of texts to process in a single API call (default: 32). Returns ------- @@ -70,8 +71,12 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: ) raise AttributeError(error_message) - response = self.client.embed(model=self.model, input=texts)["embeddings"] - return [list(inner_sequence) for inner_sequence in response] + all_embeddings = [] + for i in range(0, len(texts), batch_size): + batch = texts[i:i + batch_size] + response = self.client.embed(model=self.model, input=batch)["embeddings"] + all_embeddings.extend([list(inner_sequence) for inner_sequence in response]) + return all_embeddings async def async_get_text_embedding(self, text: str) -> List[float]: """Get embedding for a single text asynchronously.""" diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py index 135f71000..342149165 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/openai.py @@ -46,17 +46,18 @@ def get_text_embedding(self, text: str) -> List[float]: response = self.client.embeddings.create(input=text, model=self.model) return response.data[0].embedding - def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts in a single batch. + def get_texts_embeddings(self, texts: List[str], batch_size: int = 32) -> List[List[float]]: + """Get embeddings for multiple texts with automatic batch splitting. - This method efficiently processes multiple texts at once by leveraging - OpenAI's batching capabilities, which is more efficient than processing - texts individually. + This method efficiently processes multiple texts by splitting them into + smaller batches to respect API rate limits and batch size constraints. Parameters ---------- texts : List[str] A list of text strings to be embedded. + batch_size : int, optional + Maximum number of texts to process in a single API call (default: 32). Returns ------- @@ -64,11 +65,15 @@ def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: A list of embedding vectors, where each vector is a list of floats. The order of embeddings matches the order of input texts. """ - response = self.client.embeddings.create(input=texts, model=self.model) - return [data.embedding for data in response.data] + all_embeddings = [] + for i in range(0, len(texts), batch_size): + batch = texts[i:i + batch_size] + response = self.client.embeddings.create(input=batch, model=self.model) + all_embeddings.extend([data.embedding for data in response.data]) + return all_embeddings - async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts in a single batch asynchronously. + async def async_get_texts_embeddings(self, texts: List[str], batch_size: int = 32) -> List[List[float]]: + """Get embeddings for multiple texts with automatic batch splitting (async). This method should efficiently process multiple texts at once by leveraging the embedding model's batching capabilities, which is typically more efficient @@ -78,6 +83,8 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] ---------- texts : List[str] A list of text strings to be embedded. + batch_size : int, optional + Maximum number of texts to process in a single API call (default: 32). Returns ------- @@ -85,8 +92,12 @@ async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float] A list of embedding vectors, where each vector is a list of floats. The order of embeddings should match the order of input texts. """ - response = await self.aclient.embeddings.create(input=texts, model=self.model) - return [data.embedding for data in response.data] + all_embeddings = [] + for i in range(0, len(texts), batch_size): + batch = texts[i:i + batch_size] + response = await self.aclient.embeddings.create(input=batch, model=self.model) + all_embeddings.extend([data.embedding for data in response.data]) + return all_embeddings async def async_get_text_embedding(self, text: str) -> List[float]: response = await self.aclient.embeddings.create(input=[text], model=self.model) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py index ef131755d..b4a92b3a9 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py @@ -16,7 +16,7 @@ from PyCGraph import CStatus from hugegraph_llm.config import index_settings -from hugegraph_llm.models.embeddings.init_embedding import get_embedding, llm_settings +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_gremlin_example_index import ( BuildGremlinExampleIndex, @@ -37,7 +37,7 @@ def node_init(self): return CStatus(-1, "examples is required in BuildGremlinExampleIndexNode") examples = self.wk_input.examples vector_index = get_vector_index_class(index_settings.cur_vector_index) - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() self.build_gremlin_example_index_op = BuildGremlinExampleIndex( embedding, examples, vector_index 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 9e9319468..906f85407 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 @@ -14,7 +14,7 @@ # limitations under the License. from hugegraph_llm.config import index_settings -from hugegraph_llm.models.embeddings.init_embedding import get_embedding, llm_settings +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_semantic_index import BuildSemanticIndex from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -30,7 +30,7 @@ def node_init(self): from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel vector_index = get_vector_index_class(index_settings.cur_vector_index) - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() self.build_semantic_index_op = BuildSemanticIndex(embedding, vector_index) return super().node_init() diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py index 892c0879e..c073a87c9 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py @@ -14,7 +14,7 @@ # limitations under the License. from hugegraph_llm.config import index_settings -from hugegraph_llm.models.embeddings.init_embedding import get_embedding, llm_settings +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.build_vector_index import BuildVectorIndex from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -30,7 +30,7 @@ def node_init(self): from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel vector_index = get_vector_index_class(index_settings.cur_vector_index) - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() self.build_vector_index_op = BuildVectorIndex(embedding, vector_index) return super().node_init() diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py index 8881e4e9b..144b04859 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -24,7 +24,7 @@ from hugegraph_llm.operators.index_op.gremlin_example_index_query import ( GremlinExampleIndexQuery, ) -from hugegraph_llm.models.embeddings.init_embedding import get_embedding, llm_settings +from hugegraph_llm.models.embeddings.init_embedding import Embeddings class GremlinExampleIndexQueryNode(BaseNode): @@ -36,7 +36,7 @@ def node_init(self): # Build operator (index lazy-loading handled in operator) vector_index = get_vector_index_class(index_settings.cur_vector_index) - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() example_num = getattr(self.wk_input, "example_num", None) if not isinstance(example_num, int): example_num = 2 diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py index 5ab551b0f..792bfb2c3 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -18,8 +18,8 @@ from PyCGraph import CStatus from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.semantic_id_query import SemanticIdQuery -from hugegraph_llm.models.embeddings.init_embedding import get_embedding -from hugegraph_llm.config import huge_settings, index_settings, llm_settings +from hugegraph_llm.models.embeddings.init_embedding import Embeddings +from hugegraph_llm.config import huge_settings, index_settings from hugegraph_llm.utils.log import log @@ -43,7 +43,7 @@ def node_init(self): return CStatus(-1, "graph_name is required in wk_input") vector_index = get_vector_index_class(index_settings.cur_vector_index) - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() by = ( self.wk_input.semantic_by if self.wk_input.semantic_by is not None diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py index 5c6f43799..79a54fa4a 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py @@ -17,7 +17,7 @@ from hugegraph_llm.config import index_settings from hugegraph_llm.nodes.base_node import BaseNode from hugegraph_llm.operators.index_op.vector_index_query import VectorIndexQuery -from hugegraph_llm.models.embeddings.init_embedding import get_embedding, llm_settings +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.utils.log import log @@ -38,7 +38,7 @@ def node_init(self): # 从 wk_input 中读取用户配置参数 vector_index = get_vector_index_class(index_settings.cur_vector_index) - embedding = get_embedding(llm_settings) + embedding = Embeddings().get_embedding() max_items = self.wk_input.max_items if self.wk_input.max_items is not None else 3 self.operator = VectorIndexQuery(vector_index=vector_index, embedding=embedding, topk=max_items) From 90feadb079615e216792039d38c239d2f0a31fb1 Mon Sep 17 00:00:00 2001 From: LingXiao Qi Date: Mon, 27 Oct 2025 19:04:42 +0800 Subject: [PATCH 65/71] Update hugegraph-llm/src/hugegraph_llm/operators/index_op/build_semantic_index.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../hugegraph_llm/operators/index_op/build_semantic_index.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 82b8b325f..5e9e8f449 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 @@ -66,7 +66,7 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]: data.get("id_strategy") == "PRIMARY_KEY" for data in vertexlabels ) - past_vids = self.vid_index.get_all_properties() # only support Faiss + past_vids = self.vid_index.get_all_properties() # TODO: We should build vid vector index separately, especially when the vertices may be very large present_vids = context["vertices"] # Warning: data truncated by fetch_graph_data.py removed_vids = set(past_vids) - set(present_vids) From e0797fe270a05b82cde981a9a7b2e54c5743e4f4 Mon Sep 17 00:00:00 2001 From: LingXiao Qi Date: Mon, 27 Oct 2025 19:05:04 +0800 Subject: [PATCH 66/71] Update hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../hugegraph_llm/indices/vector_index/milvus_vector_store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py index b168b8c39..79e796921 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -241,7 +241,7 @@ def clean(*name: str): @staticmethod def from_name(embed_dim: int, *name: str) -> "MilvusVectorIndex": name_str = "_".join(name) - assert index_settings.milvus_host, "Qdrant host is not configured" + assert index_settings.milvus_host, "Milvus host is not configured" return MilvusVectorIndex( name_str, host=index_settings.milvus_host, From e74e09cfd3ac679fec7b6f3172136d07466be058 Mon Sep 17 00:00:00 2001 From: LingXiao Qi Date: Mon, 27 Oct 2025 19:05:27 +0800 Subject: [PATCH 67/71] Update hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py index 0b48a84dc..8a0f7ced0 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/admin_block.py @@ -69,7 +69,7 @@ def clear_llm_server_log(): # Function to validate password and control access to logs -def check_password(password, request=None): +def check_password(password, request: gr.Request | None = None): client_ip = request.client.host if request else "Unknown IP" admin_token = admin_settings.admin_token From 221effe5cee1b20c8fbd47eccb4b815383a63e9b Mon Sep 17 00:00:00 2001 From: LingXiao Qi Date: Mon, 27 Oct 2025 19:09:21 +0800 Subject: [PATCH 68/71] Update hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py | 1 - 1 file changed, 1 deletion(-) 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 8bd0f0a67..d584ccd88 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -165,7 +165,6 @@ def create_app(): # we don't need to manually check the env now # settings.check_env() prompt.update_yaml_file() - assert admin_settings.enable_login auth_enabled = admin_settings.enable_login.lower() == "true" log.info("(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled") api_auth = APIRouter(dependencies=[Depends(authenticate)] if auth_enabled else []) From b040918d39e2a8e0bd33ce0910a0ae8fa21ae23f Mon Sep 17 00:00:00 2001 From: lingxiao Date: Mon, 27 Oct 2025 19:09:35 +0800 Subject: [PATCH 69/71] fix(qdrant): use UUID for point IDs to prevent data loss from ID collisions The previous implementation used loop index (i) as point ID, which caused severe data loss when add() was called multiple times - IDs would restart from 0 and overwrite existing points. Fixed by using uuid.uuid4() to generate unique IDs for each point across all add operations, ensuring data integrity. This resolves a critical bug that could lead to inconsistent vector index and missing embeddings in Qdrant storage. --- .../indices/vector_index/qdrant_vector_store.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index 4342c90ee..613737835 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -16,6 +16,7 @@ # under the License. from typing import Any, Dict, List, Set, Union +import uuid from qdrant_client import QdrantClient from qdrant_client.http import models @@ -71,10 +72,12 @@ def add(self, vectors: List[List[float]], props: List[Any]): points = [] - for i, (vector, prop) in enumerate(zip(vectors, props)): + for vector, prop in zip(vectors, props): + # Use UUID to ensure unique point IDs across multiple add operations + # This prevents data loss from ID collisions points.append( models.PointStruct( - id=i, + id=str(uuid.uuid4()), vector=vector, payload={"property": prop}, ) From e74d11a8f1a3074cd421ef1183bc28f8a6ee4162 Mon Sep 17 00:00:00 2001 From: lingxiao Date: Mon, 27 Oct 2025 19:31:36 +0800 Subject: [PATCH 70/71] style: fix pylint issues (trailing whitespace, unused imports, line length, signature mismatch) - Remove unused 'os' import from graph_index_utils.py - Fix trailing whitespace in multiple files - Break long lines to comply with 120 char limit - Add batch_size parameter to OllamaEmbedding.async_get_texts_embeddings for consistency - Improve pylint comment placement for better readability This improves the pylint score from 9.35/10 to 9.36/10 --- .../src/hugegraph_llm/flows/get_graph_index_info.py | 4 ++-- .../hugegraph_llm/models/embeddings/init_embedding.py | 2 +- .../src/hugegraph_llm/models/embeddings/ollama.py | 5 ++++- .../nodes/index_node/build_gremlin_example_index.py | 5 +++-- .../nodes/index_node/build_semantic_index.py | 5 +++-- .../hugegraph_llm/nodes/index_node/build_vector_index.py | 5 +++-- .../nodes/index_node/gremlin_example_index_query.py | 9 ++++++--- .../nodes/index_node/semantic_id_query_node.py | 5 +++-- .../hugegraph_llm/nodes/index_node/vector_query_node.py | 5 +++-- .../src/hugegraph_llm/utils/graph_index_utils.py | 3 +-- 10 files changed, 29 insertions(+), 19 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py index 5c09eaef0..b560918ec 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py @@ -47,7 +47,7 @@ def post_deal(self, pipeline=None): from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel graph_summary_info = pipeline.getGParamWithNoEmpty("wkflow_state").to_json() - + try: vector_index_class = get_vector_index_class(index_settings.cur_vector_index) embedding = Embeddings().get_embedding() @@ -58,5 +58,5 @@ def post_deal(self, pipeline=None): except Exception: # pylint: disable=broad-except # If vector index doesn't exist or fails to load, just return graph summary pass - + return json.dumps(graph_summary_info, ensure_ascii=False, indent=2) diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py index bd903f049..26e579e41 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -101,7 +101,7 @@ def get_embedding(self): actual_dim = len(test_vec) except Exception: # pylint: disable=broad-except actual_dim = 1536 # Fallback - + # Create final instance with correct dimension embedding = LiteLLMEmbedding( embedding_dimension=actual_dim, diff --git a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py index 910372b43..28826099a 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py @@ -83,8 +83,11 @@ async def async_get_text_embedding(self, text: str) -> List[float]: response = await self.async_client.embeddings(model=self.model, prompt=text) return list(response["embedding"]) - async def async_get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: + async def async_get_texts_embeddings( + self, texts: List[str], batch_size: int = 32 + ) -> List[List[float]]: # Ollama python client may not provide batch async embeddings; fallback per item + # batch_size parameter included for consistency with base class signature results: List[List[float]] = [] for t in texts: response = await self.async_client.embeddings(model=self.model, prompt=t) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py index b4a92b3a9..90545dcdb 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_gremlin_example_index.py @@ -31,8 +31,9 @@ class BuildGremlinExampleIndexNode(BaseNode): def node_init(self): # Lazy import to avoid circular dependency - from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel - + # pylint: disable=import-outside-toplevel + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class + if not self.wk_input.examples: return CStatus(-1, "examples is required in BuildGremlinExampleIndexNode") examples = self.wk_input.examples 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 906f85407..71c853720 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 @@ -27,8 +27,9 @@ class BuildSemanticIndexNode(BaseNode): def node_init(self): # Lazy import to avoid circular dependency - from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel - + # pylint: disable=import-outside-toplevel + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class + vector_index = get_vector_index_class(index_settings.cur_vector_index) embedding = Embeddings().get_embedding() self.build_semantic_index_op = BuildSemanticIndex(embedding, vector_index) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py index c073a87c9..dfe37c6b1 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py @@ -27,8 +27,9 @@ class BuildVectorIndexNode(BaseNode): def node_init(self): # Lazy import to avoid circular dependency - from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel - + # pylint: disable=import-outside-toplevel + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class + vector_index = get_vector_index_class(index_settings.cur_vector_index) embedding = Embeddings().get_embedding() self.build_vector_index_op = BuildVectorIndex(embedding, vector_index) diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py index 144b04859..6389bbeb8 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/gremlin_example_index_query.py @@ -32,8 +32,9 @@ class GremlinExampleIndexQueryNode(BaseNode): def node_init(self): # Lazy import to avoid circular dependency - from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel - + # pylint: disable=import-outside-toplevel + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class + # Build operator (index lazy-loading handled in operator) vector_index = get_vector_index_class(index_settings.cur_vector_index) embedding = Embeddings().get_embedding() @@ -42,7 +43,9 @@ def node_init(self): example_num = 2 # Clamp to [0, 10] example_num = max(0, min(10, example_num)) - self.operator = GremlinExampleIndexQuery(vector_index=vector_index, embedding=embedding, num_examples=example_num) + self.operator = GremlinExampleIndexQuery( + vector_index=vector_index, embedding=embedding, num_examples=example_num + ) return super().node_init() def operator_schedule(self, data_json: Dict[str, Any]): diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py index 792bfb2c3..74bd9d77d 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -36,8 +36,9 @@ def node_init(self): """ try: # Lazy import to avoid circular dependency - from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel - + # pylint: disable=import-outside-toplevel + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class + graph_name = huge_settings.graph_name if not graph_name: return CStatus(-1, "graph_name is required in wk_input") diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py index 79a54fa4a..7fda64bf1 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py @@ -34,8 +34,9 @@ def node_init(self): """ try: # Lazy import to avoid circular dependency - from hugegraph_llm.utils.vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel - + # pylint: disable=import-outside-toplevel + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class + # 从 wk_input 中读取用户配置参数 vector_index = get_vector_index_class(index_settings.cur_vector_index) embedding = Embeddings().get_embedding() diff --git a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index b9b29c384..e8080b631 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py @@ -16,7 +16,6 @@ # under the License. -import os import traceback from typing import Dict, Any, Union, List @@ -44,7 +43,7 @@ def clean_all_graph_index(): # Lazy import to avoid circular dependency from .vector_index_utils import get_vector_index_class # pylint: disable=import-outside-toplevel from ..config import index_settings # pylint: disable=import-outside-toplevel - + vector_index = get_vector_index_class(index_settings.cur_vector_index) vector_index.clean(huge_settings.graph_name, "graph_vids") vector_index.clean("gremlin_examples") From 9ae2c12bbb41f6b398c5acd4983f81f50b37666c Mon Sep 17 00:00:00 2001 From: lingxiao Date: Mon, 27 Oct 2025 19:38:05 +0800 Subject: [PATCH 71/71] style: suppress remaining pylint warnings with proper justifications Fixed warnings: - W0718 (broad-exception-caught): Added pylint disable comments for legitimate broad exception handling in error recovery paths where multiple exception types need to be caught - R1711 (useless-return): Removed redundant return statement in utils.py - R0912 (too-many-branches): Suppressed for apply_vector_engine_backend() which handles multiple vector DB backends - E0401 (import-error): Suppressed for optional dependencies (pymilvus, qdrant_client) that may not be installed These changes improve pylint score from 9.36/10 to 9.38/10 while maintaining code quality and error handling robustness. --- .../src/hugegraph_llm/demo/rag_demo/configs_block.py | 4 ++-- hugegraph-llm/src/hugegraph_llm/flows/utils.py | 1 - .../hugegraph_llm/indices/vector_index/milvus_vector_store.py | 2 +- .../hugegraph_llm/indices/vector_index/qdrant_vector_store.py | 4 ++-- .../hugegraph_llm/nodes/index_node/semantic_id_query_node.py | 2 +- .../src/hugegraph_llm/nodes/index_node/vector_query_node.py | 4 ++-- 6 files changed, 8 insertions(+), 9 deletions(-) diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py index 4bf875b9c..b472ea0ba 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/configs_block.py @@ -116,7 +116,7 @@ def apply_vector_engine(engine: str): gr.Info("Configured!") -def apply_vector_engine_backend( +def apply_vector_engine_backend( # pylint: disable=too-many-branches engine: str, host: Optional[str] = None, port: Optional[str] = None, @@ -152,7 +152,7 @@ def apply_vector_engine_backend( if origin_call is None: raise gr.Error(msg) from e return -1 - except Exception as e: + except Exception as e: # pylint: disable=broad-exception-caught msg = f"Connection failed: {e}" log.error(msg) if origin_call is None: diff --git a/hugegraph-llm/src/hugegraph_llm/flows/utils.py b/hugegraph-llm/src/hugegraph_llm/flows/utils.py index 5bd7dc73c..c29c28df0 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/utils.py +++ b/hugegraph-llm/src/hugegraph_llm/flows/utils.py @@ -31,4 +31,3 @@ def prepare_schema(prepared_input: WkFlowInput, schema: str) -> None: else: log.info("Get schema '%s' from graphdb.", schema) prepared_input.graph_name = schema - return diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py index 79e796921..b83aa2836 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/milvus_vector_store.py @@ -18,7 +18,7 @@ import json from typing import Any, List, Set, Union -from pymilvus import ( +from pymilvus import ( # pylint: disable=import-error Collection, CollectionSchema, DataType, diff --git a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py index 613737835..52914d409 100644 --- a/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -18,8 +18,8 @@ from typing import Any, Dict, List, Set, Union import uuid -from qdrant_client import QdrantClient -from qdrant_client.http import models +from qdrant_client import QdrantClient # pylint: disable=import-error +from qdrant_client.http import models # pylint: disable=import-error from hugegraph_llm.config import index_settings from hugegraph_llm.indices.vector_index.base import VectorStoreBase diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py index 74bd9d77d..799b320cf 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/semantic_id_query_node.py @@ -77,7 +77,7 @@ def node_init(self): ) return super().node_init() - except Exception as e: + except Exception as e: # pylint: disable=broad-exception-caught log.error("Failed to initialize SemanticIdQueryNode: %s", e) return CStatus(-1, f"SemanticIdQueryNode initialization failed: {e}") diff --git a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py index 7fda64bf1..2ef4e8fbd 100644 --- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py +++ b/hugegraph-llm/src/hugegraph_llm/nodes/index_node/vector_query_node.py @@ -44,7 +44,7 @@ def node_init(self): self.operator = VectorIndexQuery(vector_index=vector_index, embedding=embedding, topk=max_items) return super().node_init() - except Exception as e: + except Exception as e: # pylint: disable=broad-exception-caught log.error("Failed to initialize VectorQueryNode: %s", e) from PyCGraph import CStatus @@ -73,6 +73,6 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: return data_json - except Exception as e: + except Exception as e: # pylint: disable=broad-exception-caught log.error("Vector query failed: %s", e) return data_json