This is an automated email from the ASF dual-hosted git repository.
jin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-hugegraph-ai.git
The following commit(s) were added to refs/heads/main by this push:
new 28f85757 feat(llm): support vector db layer V1.0 (#304)
28f85757 is described below
commit 28f857570f300547fec67b372797e5894df55c06
Author: LingXiao Qi <[email protected]>
AuthorDate: Tue Oct 28 19:34:31 2025 +0800
feat(llm): support vector db layer V1.0 (#304)
close #239 #246
---
## Summar
* **New Features**
* Added support for multiple vector database backends: Milvus and Qdrant
alongside Faiss.
* Introduced admin authentication system with login configuration.
* Enhanced RAG API with reranking methods, neighbor-first option, and
custom priority information.
* Added new Gremlin output types for query results.
* **Configuration**
* Improved type safety across core configurations.
* Added embedding dimension specifications for LLM models.
* New vector database configuration options.
* **Documentation**
* Updated setup instructions for vector database backends and NLTK
dependencies.
* **Refactoring**
* Redesigned vector index system to support pluggable backends.
* Restructured workflow system with extensible flow abstraction.
---
## 架构序列图
```mermaid
sequenceDiagram
participant User
participant Scheduler
participant Flow
participant GPipeline
participant Node
participant Operator
participant VectorStore
User->>Scheduler: schedule_flow(flow_name, ...)
Scheduler->>Flow: build_flow(...)
Flow->>GPipeline: 创建管道
Flow->>Node: 注册节点
Node->>Operator: 初始化操作符
Operator->>VectorStore: 获取向量存储实例
GPipeline->>Node: run()
Node->>Node: node_init()
Node->>Operator: operator_schedule(data_json)
Operator->>VectorStore: add/search/remove()
Node-->>GPipeline: 返回状态
Flow->>Flow: post_deal(pipeline)
Flow-->>Scheduler: 返回结果
Scheduler-->>User: 返回答案
```
---------
Co-authored-by: Linyu <[email protected]>
Co-authored-by: mikumifa <[email protected]>
Co-authored-by: imbajin <[email protected]>
---
hugegraph-llm/README.md | 3 +
hugegraph-llm/config.md | 13 +
hugegraph-llm/pyproject.toml | 33 ++-
hugegraph-llm/src/hugegraph_llm/config/__init__.py | 8 +-
hugegraph-llm/src/hugegraph_llm/config/generate.py | 9 +-
.../src/hugegraph_llm/config/hugegraph_config.py | 23 +-
.../{hugegraph_config.py => index_config.py} | 33 ++-
.../src/hugegraph_llm/config/models/__init__.py | 2 +
.../src/hugegraph_llm/demo/rag_demo/admin_block.py | 3 +-
.../src/hugegraph_llm/demo/rag_demo/app.py | 2 +-
.../hugegraph_llm/demo/rag_demo/configs_block.py | 127 +++++++++-
.../hugegraph_llm/flows/get_graph_index_info.py | 40 ++--
.../build_vector_index.py => flows/utils.py} | 30 +--
.../src/hugegraph_llm/indices/graph_index.py | 1 +
.../src/hugegraph_llm/indices/vector_index.py | 170 -------------
.../src/hugegraph_llm/indices/vector_index/base.py | 107 +++++++++
.../indices/vector_index/faiss_vector_store.py | 140 +++++++++++
.../indices/vector_index/milvus_vector_store.py | 263 +++++++++++++++++++++
.../indices/vector_index/qdrant_vector_store.py | 224 ++++++++++++++++++
.../src/hugegraph_llm/models/embeddings/base.py | 14 +-
.../models/embeddings/init_embedding.py | 43 +++-
.../src/hugegraph_llm/models/embeddings/litellm.py | 99 ++++++--
.../src/hugegraph_llm/models/embeddings/ollama.py | 74 ++++--
.../src/hugegraph_llm/models/embeddings/openai.py | 51 ++--
.../index_node/build_gremlin_example_index.py | 12 +-
.../nodes/index_node/build_semantic_index.py | 12 +-
.../nodes/index_node/build_vector_index.py | 12 +-
.../index_node/gremlin_example_index_query.py | 27 ++-
.../nodes/index_node/semantic_id_query_node.py | 84 ++++---
.../nodes/index_node/vector_query_node.py | 29 ++-
.../operators/common_op/check_schema.py | 2 +-
.../operators/hugegraph_op/fetch_graph_data.py | 1 -
.../operators/hugegraph_op/schema_manager.py | 2 +-
.../index_op/build_gremlin_example_index.py | 32 +--
.../operators/index_op/build_semantic_index.py | 67 +++---
.../operators/index_op/build_vector_index.py | 34 +--
.../index_op/gremlin_example_index_query.py | 69 +++---
.../operators/index_op/semantic_id_query.py | 25 +-
.../operators/index_op/vector_index_query.py | 21 +-
hugegraph-llm/src/hugegraph_llm/utils/anchor.py | 3 +-
.../src/hugegraph_llm/utils/graph_index_utils.py | 27 +--
hugegraph-llm/src/hugegraph_llm/utils/log.py | 0
.../src/hugegraph_llm/utils/vector_index_utils.py | 85 ++++---
..._vector_index.py => test_faiss_vector_index.py} | 10 +-
.../src/tests/indices/test_milvus_vector_index.py | 100 ++++++++
.../src/tests/indices/test_qdrant_vector_index.py | 102 ++++++++
hugegraph-ml/src/hugegraph_ml/models/bgrl.py | 2 +-
.../src/pyhugegraph/api/auth.py | 12 +-
.../src/pyhugegraph/api/gremlin.py | 4 +-
.../src/pyhugegraph/api/schema.py | 8 +-
.../src/pyhugegraph/api/traverser.py | 8 +-
51 files changed, 1691 insertions(+), 611 deletions(-)
diff --git a/hugegraph-llm/README.md b/hugegraph-llm/README.md
index 8b7e15c5..9e2bbe05 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 448661c9..7957ae90 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 b3894d7e..9b5f5b8d 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 =
"[email protected]" },
]
+maintainers = [
+ { name = "Apache HugeGraph Contributors", email =
"[email protected]" },
+]
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 5d4f5782..43efb0ab 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 @@ prompt.ensure_yaml_file_exists()
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 4b40e899..1bd7adea 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 @@ if __name__ == "__main__":
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 69abf0fb..1937eda1 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/hugegraph_config.py
b/hugegraph-llm/src/hugegraph_llm/config/index_config.py
similarity index 56%
copy from hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py
copy to hugegraph-llm/src/hugegraph_llm/config/index_config.py
index 69abf0fb..63895e6a 100644
--- a/hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py
+++ b/hugegraph-llm/src/hugegraph_llm/config/index_config.py
@@ -15,29 +15,24 @@
# specific language governing permissions and limitations
# under the License.
+import os
from typing import Optional
-from .models import BaseConfig
+from .models import BaseConfig
-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_space: Optional[str] = None
+class IndexConfig(BaseConfig):
+ """LLM settings"""
- # 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
+ 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
+ )
- # vector config
- vector_dis_threshold: Optional[float] = 0.9
- topk_per_keyword: Optional[int] = 1
+ 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", "")
- # rerank config
- topk_return_results: Optional[int] = 20
+ 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 e73646fd..087d8947 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 1b2032b2..8a0f7ced 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 @@ import asyncio
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 4e3f4de3..d584ccd8 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 8c595c30..b472ea0b 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 @@ import requests
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 create_configs_block() -> list:
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 86d08bf2..b560918e 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 @@ class GetGraphIndexInfoFlow(BaseFlow):
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/nodes/index_node/build_vector_index.py
b/hugegraph-llm/src/hugegraph_llm/flows/utils.py
similarity index 52%
copy from hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py
copy to hugegraph-llm/src/hugegraph_llm/flows/utils.py
index 1f6a3c75..c29c28df 100644
--- a/hugegraph-llm/src/hugegraph_llm/nodes/index_node/build_vector_index.py
+++ b/hugegraph-llm/src/hugegraph_llm/flows/utils.py
@@ -13,21 +13,21 @@
# 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.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
+import json
+from hugegraph_llm.state.ai_state import WkFlowInput
+from hugegraph_llm.utils.log import log
-class BuildVectorIndexNode(BaseNode):
- build_vector_index_op: BuildVectorIndex
- context: WkFlowState = None
- wk_input: WkFlowInput = None
- def node_init(self):
- self.build_vector_index_op =
BuildVectorIndex(get_embedding(llm_settings))
- return super().node_init()
-
- def operator_schedule(self, data_json):
- return self.build_vector_index_op.run(data_json)
+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 694ca014..31b3d21f 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 f8548318..00000000
--- 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 00000000..2e1cfc26
--- /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 00000000..a8f23155
--- /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 00000000..b83aa283
--- /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 00000000..52914d40
--- /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 698b9283..e25d61c4 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 @@ class BaseEmbedding(ABC):
"""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 @@ class BaseEmbedding(ABC):
"""
@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 de04dff8..26e579e4 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 @@ class Embeddings:
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 9d15daa0..3f9619cd 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 @@ class LiteLLMEmbedding(BaseEmbedding):
"""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 @@ class LiteLLMEmbedding(BaseEmbedding):
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 a4a8bb09..28826099 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 @@ from .base import BaseEmbedding
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 @@ class OllamaEmbedding(BaseEmbedding):
)
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 d0e15f00..34214916 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 @@ class OpenAIEmbedding:
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 @@ class OpenAIEmbedding:
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 @@ class OpenAIEmbedding:
----------
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 @@ class OpenAIEmbedding:
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 8772959d..90545dcd 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 c01cffc9..71c85372 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 1f6a3c75..dfe37c6b 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 e9283598..6389bbeb 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 typing import Any, Dict
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 68d2b72f..799b320c 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 @@ class SemanticIdQueryNode(BaseNode):
"""
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 9c8104c6..2ef4e8fb 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 @@ class VectorQueryNode(BaseNode):
"""
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 @@ class VectorQueryNode(BaseNode):
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 fc729c11..47b0f060 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 e93d916b..4c4c167c 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 @@ from pyhugegraph.client import PyHugeClient
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 90f1c00e..2f0643a7 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 @@ class SchemaManager:
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 6d9f9621..fb385a59 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 @@ class BuildGremlinExampleIndex:
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 2ed4e840..5e9e8f44 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 @@ class BuildSemanticIndex:
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 f5fb823c..64967d2d 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 b680f2ca..e3eea9f0 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 @@ class GremlinExampleIndexQuery:
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 3ac03246..49e712d0 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 @@ class SemanticIdQuery:
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 @@ class SemanticIdQuery:
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 4ed61692..5ced65b4 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 4542a7fd..a46a4e49 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 9c53e818..e8080b63 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 import FlowName
from hugegraph_llm.flows.scheduler import SchedulerSingleton
from pyhugegraph.client import PyHugeClient
-from .embedding_utils import get_filename_prefix, get_index_folder_name
from .hugegraph_utils import 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 67904a44..e96a5ade 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 0f8fd5f4..fd1eb2a1 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 00000000..b1ac0f20
--- /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 00000000..1e076805
--- /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 288a434b..0000e546 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 ab7d6616..d127c4f6 100644
--- a/hugegraph-python-client/src/pyhugegraph/api/auth.py
+++ b/hugegraph-python-client/src/pyhugegraph/api/auth.py
@@ -119,8 +119,8 @@ class AuthManager(HugeParamsBase):
@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 @@ class AuthManager(HugeParamsBase):
@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 @@ class AuthManager(HugeParamsBase):
@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 e02e7fb2..3261d60b 100644
--- a/hugegraph-python-client/src/pyhugegraph/api/gremlin.py
+++ b/hugegraph-python-client/src/pyhugegraph/api/gremlin.py
@@ -43,9 +43,7 @@ class GremlinManager(HugeParamsBase):
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 7e892667..8095887b 100644
--- a/hugegraph-python-client/src/pyhugegraph/api/schema.py
+++ b/hugegraph-python-client/src/pyhugegraph/api/schema.py
@@ -69,8 +69,8 @@ class SchemaManager(HugeParamsBase):
@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 @@ class SchemaManager(HugeParamsBase):
@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 72dddb07..2f226522 100644
--- a/hugegraph-python-client/src/pyhugegraph/api/traverser.py
+++ b/hugegraph-python-client/src/pyhugegraph/api/traverser.py
@@ -50,8 +50,8 @@ class TraverserManager(HugeParamsBase):
'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 @@ class TraverserManager(HugeParamsBase):
"&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(
