From 0f1aa621528efb41e652d8a3935f0d739614ee1a Mon Sep 17 00:00:00 2001 From: shanchunhua Date: Tue, 7 Jul 2026 20:07:01 +0800 Subject: [PATCH] feat(knowledgebase): add milvus backend support --- pyproject.toml | 2 + tests/test_milvus_knowledge_backend.py | 333 ++++++++++++++++++ veadk/config.py | 2 + veadk/configs/database_configs.py | 21 ++ .../knowledgebase/backends/milvus_backend.py | 201 +++++++++++ veadk/knowledgebase/knowledgebase.py | 17 +- 6 files changed, 574 insertions(+), 2 deletions(-) create mode 100644 tests/test_milvus_knowledge_backend.py create mode 100644 veadk/knowledgebase/backends/milvus_backend.py diff --git a/pyproject.toml b/pyproject.toml index 720dbc1f..717830a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,8 @@ extensions = [ "llama-index-llms-openai-like>=0.5.1", # For KnowledgeBase and LongTermMemory "llama-index-vector-stores-redis>=0.6.1", # For Redis database "llama-index-vector-stores-opensearch>=0.6.1", # For Opensearch database + "llama-index-vector-stores-milvus>=0.4", # For Milvus database + "pymilvus>=2.4", # For Milvus database "opensearch-py>=2.8.0", "lark-oapi", ] diff --git a/tests/test_milvus_knowledge_backend.py b/tests/test_milvus_knowledge_backend.py new file mode 100644 index 00000000..9b4de4c5 --- /dev/null +++ b/tests/test_milvus_knowledge_backend.py @@ -0,0 +1,333 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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 sys +import types + +import pytest + +from veadk.configs.database_configs import MilvusConfig +from veadk.configs.model_configs import NormalEmbeddingModelConfig +from veadk.knowledgebase.knowledgebase import _get_backend_cls + + +@pytest.fixture +def fake_milvus_dependencies(monkeypatch): + captured: dict = {} + + llama_index = types.ModuleType("llama_index") + llama_index_core = types.ModuleType("llama_index.core") + llama_index_vector_stores = types.ModuleType("llama_index.vector_stores") + llama_index_milvus = types.ModuleType("llama_index.vector_stores.milvus") + ark_embedding = types.ModuleType("veadk.models.ark_embedding") + + class FakeMilvusVectorStore: + def __init__(self, **kwargs): + captured["vector_store_kwargs"] = kwargs + self.client = FakeMilvusClient() + self.collection_name = kwargs["collection_name"] + + class FakeMilvusClient: + def get_load_state(self, collection_name): + captured.setdefault("client_events", []).append( + ("get_load_state", collection_name) + ) + return captured.get("load_state", {"state": "Loaded"}) + + def load_collection(self, collection_name): + captured.setdefault("client_events", []).append( + ("load_collection", collection_name) + ) + + class FakeStorageContext: + @classmethod + def from_defaults(cls, **kwargs): + captured["storage_context_kwargs"] = kwargs + return cls() + + class FakeNode: + text = "Milvus stores vector knowledge." + + class FakeRetriever: + def retrieve(self, query): + captured.setdefault("client_events", []).append(("retrieve", query)) + captured["query"] = query + return [FakeNode()] + + class FakeVectorStoreIndex: + def __init__(self, **kwargs): + captured["vector_index_kwargs"] = kwargs + + def as_retriever(self, similarity_top_k): + captured["similarity_top_k"] = similarity_top_k + return FakeRetriever() + + def fake_create_embedding_model(**kwargs): + captured["embedding_kwargs"] = kwargs + return "fake-embedding-model" + + llama_index_core.StorageContext = FakeStorageContext + llama_index_core.VectorStoreIndex = FakeVectorStoreIndex + llama_index_milvus.MilvusVectorStore = FakeMilvusVectorStore + ark_embedding.create_embedding_model = fake_create_embedding_model + + monkeypatch.setitem(sys.modules, "llama_index", llama_index) + monkeypatch.setitem(sys.modules, "llama_index.core", llama_index_core) + monkeypatch.setitem( + sys.modules, "llama_index.vector_stores", llama_index_vector_stores + ) + monkeypatch.setitem( + sys.modules, "llama_index.vector_stores.milvus", llama_index_milvus + ) + monkeypatch.setitem(sys.modules, "veadk.models.ark_embedding", ark_embedding) + + return captured + + +def test_get_backend_cls_returns_milvus_backend(): + from veadk.knowledgebase.backends.milvus_backend import MilvusKnowledgeBackend + + assert _get_backend_cls("milvus") is MilvusKnowledgeBackend + + +def test_milvus_config_defaults(): + config = MilvusConfig() + + assert config.uri == "" + assert config.token == "" + assert config.user == "" + assert config.password == "" + assert config.db_name == "default" + assert config.overwrite is False + assert config.timeout is None + assert config.output_fields == [] + + +def test_milvus_config_reads_environment(monkeypatch): + monkeypatch.setenv("DATABASE_MILVUS_URI", "./milvus_test.db") + monkeypatch.setenv("DATABASE_MILVUS_TOKEN", "token") + monkeypatch.setenv("DATABASE_MILVUS_DB_NAME", "kb") + monkeypatch.setenv("DATABASE_MILVUS_OVERWRITE", "true") + monkeypatch.setenv("DATABASE_MILVUS_TIMEOUT", "3.5") + monkeypatch.setenv("DATABASE_MILVUS_OUTPUT_FIELDS", "text,metadata") + + config = MilvusConfig() + + assert config.uri == "./milvus_test.db" + assert config.token == "token" + assert config.db_name == "kb" + assert config.overwrite is True + assert config.timeout == 3.5 + assert config.output_fields == "text,metadata" + + +def test_milvus_backend_initializes_vector_store(fake_milvus_dependencies): + from veadk.knowledgebase.backends.milvus_backend import MilvusKnowledgeBackend + + MilvusKnowledgeBackend( + index="company_faq", + milvus_config=MilvusConfig( + uri="./milvus.db", + user="user", + password="password", + db_name="kb", + overwrite=True, + timeout=5.0, + ), + embedding_config=NormalEmbeddingModelConfig( + name="embedding", + dim=128, + api_base="https://example.test/api/v3/", + api_key="key", + ), + ) + + assert fake_milvus_dependencies["vector_store_kwargs"] == { + "uri": "./milvus.db", + "collection_name": "company_faq", + "dim": 128, + "overwrite": True, + "token": "user:password", + "db_name": "kb", + "timeout": 5.0, + } + assert fake_milvus_dependencies["embedding_kwargs"] == { + "model_name": "embedding", + "api_key": "key", + "api_base": "https://example.test/api/v3/", + } + + +def test_milvus_backend_prefers_explicit_token(fake_milvus_dependencies): + from veadk.knowledgebase.backends.milvus_backend import MilvusKnowledgeBackend + + MilvusKnowledgeBackend( + index="company_faq", + milvus_config=MilvusConfig( + uri="./milvus.db", + token="explicit", + user="user", + password="password", + ), + embedding_config=NormalEmbeddingModelConfig( + name="embedding", + dim=128, + api_base="https://example.test/api/v3/", + api_key="key", + ), + ) + + assert fake_milvus_dependencies["vector_store_kwargs"]["token"] == "explicit" + + +def test_milvus_backend_passes_output_fields(fake_milvus_dependencies): + from veadk.knowledgebase.backends.milvus_backend import MilvusKnowledgeBackend + + MilvusKnowledgeBackend( + index="company_faq", + milvus_config=MilvusConfig( + uri="./milvus.db", + output_fields=["text", "metadata"], + ), + embedding_config=NormalEmbeddingModelConfig( + name="embedding", + dim=128, + api_base="https://example.test/api/v3/", + api_key="key", + ), + ) + + assert fake_milvus_dependencies["vector_store_kwargs"]["output_fields"] == [ + "text", + "metadata", + ] + + +def test_milvus_backend_parses_output_fields_string(fake_milvus_dependencies): + from veadk.knowledgebase.backends.milvus_backend import MilvusKnowledgeBackend + + MilvusKnowledgeBackend( + index="company_faq", + milvus_config=MilvusConfig( + uri="./milvus.db", + output_fields="text, metadata", + ), + embedding_config=NormalEmbeddingModelConfig( + name="embedding", + dim=128, + api_base="https://example.test/api/v3/", + api_key="key", + ), + ) + + assert fake_milvus_dependencies["vector_store_kwargs"]["output_fields"] == [ + "text", + "metadata", + ] + + +@pytest.mark.parametrize("index", ["", "1bad", "bad-name", "bad.name"]) +def test_milvus_backend_rejects_invalid_collection_names(monkeypatch, index): + from veadk.knowledgebase.backends.milvus_backend import MilvusKnowledgeBackend + + monkeypatch.setattr(MilvusKnowledgeBackend, "model_post_init", lambda *_: None) + backend = MilvusKnowledgeBackend(index=index) + + with pytest.raises(ValueError, match="Milvus collection name"): + backend.precheck_index_naming() + + +def test_milvus_backend_requires_uri(fake_milvus_dependencies): + from veadk.knowledgebase.backends.milvus_backend import MilvusKnowledgeBackend + + with pytest.raises(ValueError, match="Milvus uri must be configured"): + MilvusKnowledgeBackend( + index="company_faq", + embedding_config=NormalEmbeddingModelConfig( + name="embedding", + dim=128, + api_base="https://example.test/api/v3/", + api_key="key", + ), + ) + + +def test_knowledgebase_milvus_search_wraps_strings( + monkeypatch, fake_milvus_dependencies +): + monkeypatch.setenv("MODEL_EMBEDDING_API_KEY", "key") + monkeypatch.setenv("DATABASE_MILVUS_URI", "./milvus.db") + + from veadk.knowledgebase import KnowledgeBase + from veadk.knowledgebase.backends.milvus_backend import MilvusKnowledgeBackend + + kb = KnowledgeBase(backend="milvus", index="company_faq") + + assert isinstance(kb._backend, MilvusKnowledgeBackend) + entries = kb.search("what is Milvus?", top_k=3) + + assert entries[0].content == "Milvus stores vector knowledge." + assert fake_milvus_dependencies["query"] == "what is Milvus?" + assert fake_milvus_dependencies["similarity_top_k"] == 3 + + +def test_milvus_backend_loads_released_collection_before_search( + fake_milvus_dependencies, +): + fake_milvus_dependencies["load_state"] = {"state": "released"} + + from veadk.knowledgebase.backends.milvus_backend import MilvusKnowledgeBackend + + backend = MilvusKnowledgeBackend( + index="company_faq", + milvus_config=MilvusConfig(uri="./milvus.db"), + embedding_config=NormalEmbeddingModelConfig( + name="embedding", + dim=128, + api_base="https://example.test/api/v3/", + api_key="key", + ), + ) + + backend.search("what is Milvus?", top_k=3) + + assert fake_milvus_dependencies["client_events"] == [ + ("get_load_state", "company_faq"), + ("load_collection", "company_faq"), + ("retrieve", "what is Milvus?"), + ] + + +def test_milvus_backend_does_not_load_loaded_collection(fake_milvus_dependencies): + fake_milvus_dependencies["load_state"] = {"state": "Loaded"} + + from veadk.knowledgebase.backends.milvus_backend import MilvusKnowledgeBackend + + backend = MilvusKnowledgeBackend( + index="company_faq", + milvus_config=MilvusConfig(uri="./milvus.db"), + embedding_config=NormalEmbeddingModelConfig( + name="embedding", + dim=128, + api_base="https://example.test/api/v3/", + api_key="key", + ), + ) + + backend.search("what is Milvus?", top_k=3) + + assert fake_milvus_dependencies["client_events"] == [ + ("get_load_state", "company_faq"), + ("retrieve", "what is Milvus?"), + ] diff --git a/veadk/config.py b/veadk/config.py index 3ad0cc4c..a1c4ce9e 100644 --- a/veadk/config.py +++ b/veadk/config.py @@ -21,6 +21,7 @@ from veadk.configs.auth_configs import VeIdentityConfig from veadk.configs.model_configs import RealtimeModelConfig from veadk.configs.database_configs import ( + MilvusConfig, MysqlConfig, OpensearchConfig, RedisConfig, @@ -79,6 +80,7 @@ class VeADKConfig(BaseModel): opensearch: OpensearchConfig = Field(default_factory=OpensearchConfig) mysql: MysqlConfig = Field(default_factory=MysqlConfig) redis: RedisConfig = Field(default_factory=RedisConfig) + milvus: MilvusConfig = Field(default_factory=MilvusConfig) viking_knowledgebase: VikingKnowledgebaseConfig = Field( default_factory=VikingKnowledgebaseConfig ) diff --git a/veadk/configs/database_configs.py b/veadk/configs/database_configs.py index ed8c09f7..e4f7bec0 100644 --- a/veadk/configs/database_configs.py +++ b/veadk/configs/database_configs.py @@ -15,6 +15,7 @@ import os from functools import cached_property +from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict from veadk.consts import DEFAULT_TOS_BUCKET_NAME @@ -89,6 +90,26 @@ class RedisConfig(BaseSettings): """STS token for Redis auth, not supported yet.""" +class MilvusConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="DATABASE_MILVUS_") + + uri: str = "" + + token: str = "" + + user: str = "" + + password: str = "" + + db_name: str = "default" + + overwrite: bool = False + + timeout: float | None = None + + output_fields: list[str] | str = Field(default_factory=list) + + class Mem0Config(BaseSettings): model_config = SettingsConfigDict(env_prefix="DATABASE_MEM0_") diff --git a/veadk/knowledgebase/backends/milvus_backend.py b/veadk/knowledgebase/backends/milvus_backend.py new file mode 100644 index 00000000..13fc2676 --- /dev/null +++ b/veadk/knowledgebase/backends/milvus_backend.py @@ -0,0 +1,201 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import re +from typing import Any + +from pydantic import Field +from typing_extensions import override + +import veadk.config # noqa E401 +from veadk.configs.database_configs import MilvusConfig +from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig +from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend + + +class MilvusKnowledgeBackend(BaseKnowledgebaseBackend): + """Milvus-based backend for knowledgebase. + + Milvus backend stores embedded chunks in a Milvus collection through + LlamaIndex's Milvus vector store. ``index`` maps directly to the Milvus + collection name. + """ + + milvus_config: MilvusConfig = Field(default_factory=MilvusConfig) + """Milvus connection config.""" + + embedding_config: EmbeddingModelConfig | NormalEmbeddingModelConfig = Field( + default_factory=EmbeddingModelConfig + ) + """Embedding model configs.""" + + def model_post_init(self, __context: Any) -> None: + self.precheck_index_naming() + self._precheck_milvus_uri() + + try: + from llama_index.core import StorageContext, VectorStoreIndex + from llama_index.vector_stores.milvus import MilvusVectorStore + from veadk.models.ark_embedding import create_embedding_model + except ImportError as e: + raise ImportError( + "Please install VeADK extensions\npip install veadk-python[extensions]" + ) from e + + self._embed_model = create_embedding_model( + model_name=self.embedding_config.name, + api_key=self.embedding_config.api_key, + api_base=self.embedding_config.api_base, + ) + + vector_store_kwargs: dict[str, Any] = { + "uri": self.milvus_config.uri, + "collection_name": self.index, + "dim": self.embedding_config.dim, + "overwrite": self.milvus_config.overwrite, + } + + token = self._resolve_token() + if token: + vector_store_kwargs["token"] = token + + if self.milvus_config.db_name: + vector_store_kwargs["db_name"] = self.milvus_config.db_name + + if self.milvus_config.timeout is not None: + vector_store_kwargs["timeout"] = self.milvus_config.timeout + + output_fields = self._resolve_output_fields() + if output_fields: + vector_store_kwargs["output_fields"] = output_fields + + self._vector_store = MilvusVectorStore(**vector_store_kwargs) + self._storage_context = StorageContext.from_defaults( + vector_store=self._vector_store + ) + self._vector_index = VectorStoreIndex( + nodes=[], + storage_context=self._storage_context, + embed_model=self._embed_model, + ) + + @override + def precheck_index_naming(self) -> None: + if not isinstance(self.index, str) or not self.index: + raise ValueError("Milvus collection name must not be empty.") + if len(self.index) > 255: + raise ValueError("Milvus collection name is too long.") + if not re.fullmatch(r"^[A-Za-z_][A-Za-z0-9_]*$", self.index): + raise ValueError( + "Milvus collection name must start with a letter or underscore " + "and contain only letters, numbers, and underscores." + ) + + def _precheck_milvus_uri(self) -> None: + if not self.milvus_config.uri: + raise ValueError( + "Milvus uri must be configured via DATABASE_MILVUS_URI or " + "MilvusConfig(uri=...)." + ) + + @override + def add_from_directory(self, directory: str) -> bool: + from llama_index.core import SimpleDirectoryReader + + documents = SimpleDirectoryReader(input_dir=directory).load_data() + nodes = self._split_documents(documents) + self._vector_index.insert_nodes(nodes) + return True + + @override + def add_from_files(self, files: list[str]) -> bool: + from llama_index.core import SimpleDirectoryReader + + documents = SimpleDirectoryReader(input_files=files).load_data() + nodes = self._split_documents(documents) + self._vector_index.insert_nodes(nodes) + return True + + @override + def add_from_text(self, text: str | list[str]) -> bool: + from llama_index.core import Document + + if isinstance(text, str): + documents = [Document(text=text)] + else: + documents = [Document(text=t) for t in text] + nodes = self._split_documents(documents) + self._vector_index.insert_nodes(nodes) + return True + + @override + def search(self, query: str, top_k: int = 5) -> list[str]: + self._ensure_collection_loaded() + _retriever = self._vector_index.as_retriever(similarity_top_k=top_k) + retrieved_nodes = _retriever.retrieve(query) + return [node.text for node in retrieved_nodes] + + def _resolve_token(self) -> str | None: + if self.milvus_config.token: + return self.milvus_config.token + if self.milvus_config.user and self.milvus_config.password: + return f"{self.milvus_config.user}:{self.milvus_config.password}" + return None + + def _resolve_output_fields(self) -> list[str]: + output_fields = self.milvus_config.output_fields + if isinstance(output_fields, str): + output_fields = output_fields.strip() + if not output_fields: + return [] + if output_fields.startswith("["): + parsed_output_fields = json.loads(output_fields) + if not isinstance(parsed_output_fields, list) or not all( + isinstance(field, str) for field in parsed_output_fields + ): + raise ValueError("Milvus output_fields must be a list of strings.") + output_fields = parsed_output_fields + else: + output_fields = output_fields.split(",") + + return [field.strip() for field in output_fields if field.strip()] + + def _ensure_collection_loaded(self) -> None: + vector_store = getattr(self, "_vector_store", None) + client = getattr(vector_store, "client", None) + collection_name = getattr(vector_store, "collection_name", self.index) + if client is None or not collection_name: + return + + try: + load_state = client.get_load_state(collection_name=collection_name) + except Exception: + client.load_collection(collection_name=collection_name) + return + + state = load_state.get("state") if isinstance(load_state, dict) else load_state + if "loaded" not in str(state).lower(): + client.load_collection(collection_name=collection_name) + + def _split_documents(self, documents: list[Any]) -> list[Any]: + """Split document into chunks.""" + from veadk.knowledgebase.backends.utils import get_llama_index_splitter + + nodes = [] + for document in documents: + splitter = get_llama_index_splitter(document.metadata.get("file_path", "")) + _nodes = splitter.get_nodes_from_documents([document]) + nodes.extend(_nodes) + return nodes diff --git a/veadk/knowledgebase/knowledgebase.py b/veadk/knowledgebase/knowledgebase.py index b0b110d1..93c8eeca 100644 --- a/veadk/knowledgebase/knowledgebase.py +++ b/veadk/knowledgebase/knowledgebase.py @@ -47,6 +47,12 @@ def _get_backend_cls(backend: str) -> type[BaseKnowledgebaseBackend]: ) return RedisKnowledgeBackend + case "milvus": + from veadk.knowledgebase.backends.milvus_backend import ( + MilvusKnowledgeBackend, + ) + + return MilvusKnowledgeBackend case "tos_vector": from veadk.knowledgebase.backends.tos_vector_backend import ( TosVectorKnowledgeBackend, @@ -87,12 +93,13 @@ class KnowledgeBase(BaseModel): Attributes: name (str): The name of the knowledge base. Default is "user_knowledgebase". description (str): A description of the knowledge base. Default is "This knowledgebase stores some user-related information." - backend (Union[Literal["local", "opensearch", "viking", "redis"], BaseKnowledgebaseBackend]): + backend (Union[Literal["local", "opensearch", "viking", "redis", "milvus"], BaseKnowledgebaseBackend]): The type of backend to use for storing and querying the knowledge base. Supported options include: - 'local' for in-memory storage (data is lost when the program exits). - 'opensearch' for OpenSearch (requires OpenSearch cluster). - 'viking' for Volcengine VikingDB (requires VikingDB service). - 'redis' for Redis with vector search capability (requires Redis). + - 'milvus' for Milvus vector database (requires Milvus). Default is 'local'. backend_config (dict): Configuration dictionary for the selected backend. top_k (int): The number of top similar documents to retrieve during a search. Default is 10. @@ -109,7 +116,13 @@ class KnowledgeBase(BaseModel): backend: Union[ Literal[ - "local", "opensearch", "viking", "redis", "tos_vector", "context_search" + "local", + "opensearch", + "viking", + "redis", + "milvus", + "tos_vector", + "context_search", ], BaseKnowledgebaseBackend, ] = "local"