diff --git a/hugegraph-llm/README.md b/hugegraph-llm/README.md index 8b7e15c50..9e2bbe05d 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 448661c96..7957ae904 100644 --- a/hugegraph-llm/config.md +++ b/hugegraph-llm/config.md @@ -16,6 +16,7 @@ - [LiteLLM 配置](#litellm-配置) - [重排序配置](#重排序配置) - [HugeGraph 数据库配置](#hugegraph-数据库配置) + - [向量数据库配置](#向量数据库配置) - [管理员配置](#管理员配置) - [配置使用示例](#配置使用示例) - [配置文件位置](#配置文件位置) @@ -130,6 +131,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 b3894d7e5..9b5f5b8d0 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,9 +40,9 @@ dependencies = [ "numpy", "pandas", "pydantic", + "tqdm", "scipy", "python-igraph", - # LLM specific dependencies "openai", @@ -63,6 +64,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" @@ -92,3 +100,22 @@ hugegraph-python-client = { workspace = true } # 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"} + +[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/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/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..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 @@ -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: gr.Request | None = 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 4e3f4de39..d584ccd88 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py +++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py @@ -116,7 +116,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, 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..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 @@ -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( # pylint: disable=too-many-branches + 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: # pylint: disable=broad-exception-caught + 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/flows/get_graph_index_info.py b/hugegraph-llm/src/hugegraph_llm/flows/get_graph_index_info.py index 86d08bf2d..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 @@ -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) 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..c29c28df0 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/flows/utils.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. + +import json + +from hugegraph_llm.state.ai_state import WkFlowInput +from hugegraph_llm.utils.log import log + + +def prepare_schema(prepared_input: WkFlowInput, schema: str) -> None: + 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 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..b83aa2836 --- /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 ( # pylint: disable=import-error + 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, "Milvus 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..52914d409 --- /dev/null +++ b/hugegraph-llm/src/hugegraph_llm/indices/vector_index/qdrant_vector_store.py @@ -0,0 +1,224 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from typing import Any, Dict, List, Set, Union +import uuid + +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 +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 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=str(uuid.uuid4()), + 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/models/embeddings/base.py b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py index 698b92837..e25d61c4f 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/base.py @@ -61,8 +61,14 @@ def get_text_embedding(self, text: str) -> List[float]: """Comment""" @abstractmethod - def get_texts_embeddings(self, texts: List[str]) -> List[List[float]]: - """Get embeddings for multiple texts in a single batch. + def get_embedding_dim( + self, + ) -> int: + """Get the dimension of the embedding.""" + + @abstractmethod + 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 @@ -81,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 de04dff87..26e579e41 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/init_embedding.py @@ -57,23 +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": - return OpenAIEmbedding( + # 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": - return OllamaEmbedding( - model_name=llm_settings.ollama_embedding_model, + # 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( + # 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, ) + # 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 9d15daa0a..3f9619cd9 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, @@ -62,32 +64,81 @@ 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_name, - 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 - 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, - input=texts, + model=self.model, + input=text, 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] + 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], 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: + 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 a4a8bb098..28826099a 100644 --- a/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py +++ b/hugegraph-llm/src/hugegraph_llm/models/embeddings/ollama.py @@ -23,18 +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], batch_size: int = 32) -> List[List[float]]: + """Get embeddings for multiple texts with automatic batch splitting. - 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 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 ------- @@ -49,23 +71,25 @@ 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"] - 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_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], 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) + 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..342149165 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,24 +33,31 @@ 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]]: - """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 ------- @@ -56,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_name) - 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 @@ -70,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 ------- @@ -77,5 +92,13 @@ 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) - 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) + return response.data[0].embedding 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..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 @@ -15,8 +15,8 @@ from PyCGraph import CStatus -from hugegraph_llm.config import llm_settings -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +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 ( BuildGremlinExampleIndex, @@ -30,12 +30,18 @@ class BuildGremlinExampleIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): + # Lazy import to avoid circular dependency + # 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 + vector_index = get_vector_index_class(index_settings.cur_vector_index) + embedding = Embeddings().get_embedding() self.build_gremlin_example_index_op = BuildGremlinExampleIndex( - get_embedding(llm_settings), examples + embedding, examples, vector_index ) return super().node_init() 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..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 @@ -13,8 +13,8 @@ # 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.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 @@ -26,7 +26,13 @@ class BuildSemanticIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): - self.build_semantic_index_op = BuildSemanticIndex(get_embedding(llm_settings)) + # Lazy import to avoid circular dependency + # 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) 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..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 @@ -13,8 +13,8 @@ # 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.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 from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState @@ -26,7 +26,13 @@ class BuildVectorIndexNode(BaseNode): wk_input: WkFlowInput = None def node_init(self): - self.build_vector_index_op = BuildVectorIndex(get_embedding(llm_settings)) + # Lazy import to avoid circular dependency + # 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) 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..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 @@ -19,36 +19,37 @@ from PyCGraph import CStatus -from hugegraph_llm.config import llm_settings +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, ) -from hugegraph_llm.models.embeddings.init_embedding import get_embedding +from hugegraph_llm.models.embeddings.init_embedding import Embeddings class GremlinExampleIndexQueryNode(BaseNode): operator: GremlinExampleIndexQuery def node_init(self): + # Lazy import to avoid circular dependency + # 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) - embedding = get_embedding(llm_settings) + 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 + vector_index=vector_index, embedding=embedding, num_examples=example_num ) - return CStatus() + return super().node_init() 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) + try: + return self.operator.run(data_json) + except ValueError as err: + return {"status": CStatus(-1, str(err))} 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 68d2b72f2..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 @@ -14,11 +14,12 @@ # 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.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, index_settings from hugegraph_llm.utils.log import log @@ -33,42 +34,53 @@ def node_init(self): """ Initialize the semantic ID query operator. """ - 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 - 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 - ) + try: + # Lazy import to avoid circular dependency + # pylint: disable=import-outside-toplevel + from hugegraph_llm.utils.vector_index_utils import get_vector_index_class - # 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, - ) + 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 + 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 + ) + + # 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, + vector_dis_threshold=vector_dis_threshold, + ) + + return super().node_init() + except Exception as e: # pylint: disable=broad-exception-caught + log.error("Failed to initialize SemanticIdQueryNode: %s", e) - return super().node_init() + return CStatus(-1, f"SemanticIdQueryNode initialization failed: {e}") def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: """ 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 9c8104c6e..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 @@ -14,10 +14,10 @@ # limitations under the License. from typing import Dict, Any -from hugegraph_llm.config import llm_settings +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 +from hugegraph_llm.models.embeddings.init_embedding import Embeddings from hugegraph_llm.utils.log import log @@ -32,14 +32,23 @@ def node_init(self): """ Initialize the vector query operator """ - # 从 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 - ) + try: + # Lazy import to avoid circular dependency + # 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() + 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) + return super().node_init() + except Exception as e: # pylint: disable=broad-exception-caught + log.error("Failed to initialize VectorQueryNode: %s", e) + from PyCGraph import CStatus - self.operator = VectorIndexQuery(embedding=embedding, topk=max_items) - return super().node_init() + return CStatus(-1, f"VectorQueryNode initialization failed: {e}") def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: """ @@ -64,6 +73,6 @@ def operator_schedule(self, data_json: Dict[str, Any]) -> Dict[str, Any]: return data_json - except ValueError as e: + except Exception as e: # pylint: disable=broad-exception-caught log.error("Vector query failed: %s", e) return 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 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..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 @@ -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() # 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/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/graph_index_utils.py b/hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py index 9c53e8183..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 @@ -25,13 +24,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 +40,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!") 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 67904a445..e96a5adec 100644 --- a/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py +++ b/hugegraph-llm/src/hugegraph_llm/utils/vector_index_utils.py @@ -16,20 +16,16 @@ # 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.flows import FlowName -from hugegraph_llm.indices.vector_index import VectorIndex -from hugegraph_llm.models.embeddings.init_embedding import model_map +from hugegraph_llm.config import huge_settings, index_settings from hugegraph_llm.flows.scheduler import SchedulerSingleton -from hugegraph_llm.utils.embedding_utils import ( - get_filename_prefix, - get_index_folder_name, -) +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 def read_documents(input_file, input_text): @@ -63,28 +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, @@ -92,15 +75,8 @@ 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!") @@ -109,4 +85,37 @@ 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(FlowName.BUILD_VECTOR_INDEX, texts) + return scheduler.schedule_flow("build_vector_index", texts) + + +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/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(