refine(kb): reconstruct knowledgebase

Author: yaozheng-fang

veadk/configs/model_configs.py

@@ -40,3 +40,20 @@ class ModelConfig(BaseSettings):
     @cached_property
     def api_key(self) -> str:
         return os.getenv("MODEL_AGENT_API_KEY") or ARKVeAuth().token
+
+
+class EmbeddingModelConfig(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="MODEL_EMBEDDING_")
+
+    name: str = "doubao-embedding-text-240715"
+    """Model name for embedding."""
+
+    dim: int = 2560
+    """Embedding dim is different from different models."""
+
+    api_base: str = "https://ark.cn-beijing.volces.com/api/v3/embeddings"
+    """The api base of the model for embedding."""
+
+    @cached_property
+    def api_key(self) -> str:
+        return os.getenv("MODEL_EMBEDDING_API_KEY") or ARKVeAuth().token

veadk/knowledgebase/backends/__init__.py

@@ -0,0 +1,13 @@
+# 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.

veadk/knowledgebase/backends/base_backend.py

@@ -0,0 +1,54 @@
+# 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.
+
+from abc import ABC, abstractmethod
+
+from pydantic import BaseModel
+
+
+class BaseKnowledgebaseBackend(ABC, BaseModel):
+    index: str
+    """Index or collection name of the vector storage."""
+
+    @abstractmethod
+    def add_from_directory(self, directory: str, **kwargs) -> bool:
+        """Add knowledge from file path to knowledgebase"""
+        ...
+
+    @abstractmethod
+    def add_from_files(self, files: list[str], **kwargs) -> bool:
+        """Add knowledge (e.g, documents, strings, ...) to knowledgebase"""
+        ...
+
+    @abstractmethod
+    def add_from_text(self, text: str | list[str], **kwargs) -> bool:
+        """Add knowledge from text to knowledgebase"""
+        ...
+
+    @abstractmethod
+    def search(self, **kwargs) -> list:
+        """Search knowledge from knowledgebase"""
+        ...
+
+    def delete(self, **kwargs) -> bool:
+        """Delete knowledge from knowledgebase"""
+        ...
+
+    def list_docs(self, **kwargs) -> None:
+        """List original documents in knowledgebase"""
+        pass
+
+    def list_chunks(self, **kwargs) -> None:
+        """List embeded document chunks in knowledgebase"""
+        pass

veadk/knowledgebase/backends/in_memory_backend.py

@@ -0,0 +1,75 @@
+# 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.
+
+from llama_index.core import Document, SimpleDirectoryReader, VectorStoreIndex
+from llama_index.core.schema import BaseNode
+from llama_index.embeddings.openai_like import OpenAILikeEmbedding
+from pydantic import Field
+from typing_extensions import Any, override
+
+from veadk.configs.model_configs import EmbeddingModelConfig
+from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
+from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+
+
+class InMemoryKnowledgeBackend(BaseKnowledgebaseBackend):
+    embedding_config: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
+    """Embedding model configs"""
+
+    def model_post_init(self, __context: Any) -> None:
+        self._embed_model = OpenAILikeEmbedding(
+            model_name=self.embedding_config.name,
+            api_key=self.embedding_config.api_key,
+            api_base=self.embedding_config.api_base,
+        )
+        self._vector_index = VectorStoreIndex([], embed_model=self._embed_model)
+        self._retriever = self._vector_index.as_retriever()
+
+    @override
+    def add_from_directory(self, directory: str) -> bool:
+        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:
+        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:
+        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]:
+        retrieved_nodes = self._retriever.retrieve(query, top_k=top_k)
+        return [node.text for node in retrieved_nodes]
+
+    def _split_documents(self, documents: list[Document]) -> list[BaseNode]:
+        """Split document into chunks"""
+        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

veadk/knowledgebase/backends/opensearch_backend.py

@@ -0,0 +1,112 @@
+# 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.
+
+from llama_index.core import (
+    Document,
+    SimpleDirectoryReader,
+    StorageContext,
+    VectorStoreIndex,
+)
+from llama_index.core.schema import BaseNode
+from llama_index.embeddings.openai_like import OpenAILikeEmbedding
+from llama_index.vector_stores.opensearch import (
+    OpensearchVectorClient,
+    OpensearchVectorStore,
+)
+from pydantic import Field
+from typing_extensions import Any, override
+
+from veadk.configs.database_configs import OpensearchConfig
+from veadk.configs.model_configs import EmbeddingModelConfig
+from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
+from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+
+
+class OpensearchKnowledgeBackend(BaseKnowledgebaseBackend):
+    opensearch_config: OpensearchConfig = Field(default_factory=OpensearchConfig)
+    """Opensearch client configs"""
+
+    embedding_config: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
+    """Embedding model configs"""
+
+    def model_post_init(self, __context: Any) -> None:
+        self._opensearch_client = OpensearchVectorClient(
+            endpoint=self.opensearch_config.host,
+            port=self.opensearch_config.port,
+            http_auth=(
+                self.opensearch_config.username,
+                self.opensearch_config.password,
+            ),
+            use_ssl=True,
+            verify_certs=False,
+            dim=self.embedding_config.dim,
+            index=self.index,  # collection name
+        )
+
+        self._vector_store = OpensearchVectorStore(client=self._opensearch_client)
+
+        self._storage_context = StorageContext.from_defaults(
+            vector_store=self._vector_store
+        )
+
+        self._embed_model = OpenAILikeEmbedding(
+            model_name=self.embedding_config.name,
+            api_key=self.embedding_config.api_key,
+            api_base=self.embedding_config.api_base,
+        )
+
+        self._vector_index = VectorStoreIndex.from_documents(
+            documents=[],
+            storage_context=self._storage_context,
+            embed_model=self._embed_model,
+        )
+        self._retriever = self._vector_index.as_retriever()
+
+    @override
+    def add_from_directory(self, directory: str) -> bool:
+        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:
+        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:
+        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]:
+        retrieved_nodes = self._retriever.retrieve(query, top_k=top_k)
+        return [node.text for node in retrieved_nodes]
+
+    def _split_documents(self, documents: list[Document]) -> list[BaseNode]:
+        """Split document into chunks"""
+        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

veadk/knowledgebase/backends/redis_backend.py

@@ -0,0 +1,108 @@
+# 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.
+
+from llama_index.core import (
+    Document,
+    SimpleDirectoryReader,
+    StorageContext,
+    VectorStoreIndex,
+)
+from llama_index.core.schema import BaseNode
+from llama_index.embeddings.openai_like import OpenAILikeEmbedding
+from llama_index.vector_stores.redis import RedisVectorStore
+from pydantic import Field
+from redis import Redis
+from typing_extensions import Any, override
+
+from veadk.configs.database_configs import RedisConfig
+from veadk.configs.model_configs import EmbeddingModelConfig
+from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
+from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+
+
+class RedisKnowledgeBackend(BaseKnowledgebaseBackend):
+    redis_config: RedisConfig = Field(default_factory=RedisConfig)
+    """Redis client configs"""
+
+    embedding_config: EmbeddingModelConfig
+    """Embedding model configs"""
+
+    def model_post_init(self, __context: Any) -> None:
+        # We will use `from_url` to init Redis client once the
+        # AK/SK -> STS token is ready.
+        # self._redis_client = Redis.from_url(url=...)
+
+        self._redis_client = Redis(
+            host=self.redis_config.host,
+            port=self.redis_config.port,
+            db=self.redis_config.db,
+            password=self.redis_config.password,
+        )
+
+        self._embed_model = OpenAILikeEmbedding(
+            model_name=self.embedding_config.name,
+            api_key=self.embedding_config.api_key,
+            api_base=self.embedding_config.api_base,
+        )
+
+        self._vector_store = RedisVectorStore(
+            redis_client=self._redis_client, overwrite=True
+        )
+
+        self._storage_context = StorageContext.from_defaults(
+            vector_store=self._vector_store
+        )
+
+        self._vector_index = VectorStoreIndex.from_documents(
+            documents=[], storage_context=self._storage_context
+        )
+        self._retriever = self._vector_index.as_retriever()
+
+    @override
+    def add_from_directory(self, directory: str) -> bool:
+        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:
+        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:
+        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]:
+        retrieved_nodes = self._retriever.retrieve(query, top_k=top_k)
+        return [node.text for node in retrieved_nodes]
+
+    def _split_documents(self, documents: list[Document]) -> list[BaseNode]:
+        """Split document into chunks"""
+        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