Skip to content

Commit 96efdbc

Browse files
zakahanzakahan
authored
feat(kb): tos vector backend support (#369)
* feat: tos_backend * feat: tos for knowledgebase * feat: add tos base * fix: region * fix: tos logger and config * fix: one line
1 parent 1ce1662 commit 96efdbc

File tree

4 files changed

+247
-1
lines changed

4 files changed

+247
-1
lines changed

‎config.yaml.full‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,11 @@ database:
142142
mem0:
143143
base_url: # default "https://api.mem0.ai/v1", using full https url including port
144144
api_key: #api_key
145+
tos_vector:
146+
endpoint: tosvectors-cn-beijing.volces.com # default Volcengine TOS Vector endpoint
147+
region: cn-beijing # default Volcengine TOS Vector region
148+
bucket:
149+
account_id:
145150

146151

147152

‎veadk/configs/database_configs.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,3 +130,47 @@ class NormalTOSConfig(BaseSettings):
130130
region: str = "cn-beijing"
131131

132132
bucket: str
133+
134+
135+
class TOSVectorConfig(BaseSettings):
136+
model_config = SettingsConfigDict(env_prefix="DATABASE_TOS_VECTOR_")
137+
138+
endpoint: str = "tosvectors-cn-beijing.volces.com"
139+
140+
region: str = "cn-beijing"
141+
142+
security_token: str | None = None
143+
144+
max_retry_count: int = 3
145+
146+
max_connections: int = 1024
147+
148+
connection_time: int = 10
149+
150+
enable_verify_ssl: bool = True
151+
152+
dns_cache_time: int = 15
153+
154+
proxy_host: str | None = None
155+
156+
proxy_port: int | None = None
157+
158+
proxy_username: str | None = None
159+
160+
proxy_password: str | None = None
161+
162+
high_latency_log_threshold: int = 100
163+
164+
socket_timeout: int = 30
165+
166+
credentials_provider: object | None = None
167+
168+
except100_continue_threshold: int = 65536
169+
170+
user_agent_product_name: str | None = None
171+
172+
user_agent_soft_name: str | None = None
173+
174+
user_agent_soft_version: str | None = None
175+
176+
user_agent_customized_key_values: dict[str, str] | None = None

‎veadk/knowledgebase/backends/tos_vector_backend.py‎

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import json
16+
import os
17+
18+
from llama_index.core import (
19+
Document,
20+
SimpleDirectoryReader,
21+
)
22+
from llama_index.core.schema import BaseNode
23+
from llama_index.embeddings.openai_like import OpenAILikeEmbedding
24+
from pydantic import Field
25+
from typing_extensions import Any, override
26+
27+
import veadk.config # noqa E401
28+
from veadk.configs.database_configs import TOSVectorConfig
29+
from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
30+
31+
from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
32+
from veadk.knowledgebase.backends.utils import get_llama_index_splitter
33+
from veadk.utils.logger import get_logger
34+
35+
logger = get_logger(__name__)
36+
try:
37+
from tos.vector_client import VectorClient
38+
from tos import DataType, DistanceMetricType
39+
from tos.exceptions import TosServerError
40+
from tos.models2 import Vector, VectorData
41+
except ImportError:
42+
raise ImportError(
43+
"Please install VeADK extensions\npip install veadk-python[extensions]"
44+
)
45+
46+
47+
class TosVectorKnowledgeBackend(BaseKnowledgebaseBackend):
48+
"""TOS-based backend for knowledgebase."""
49+
50+
volcengine_access_key: str | None = Field(
51+
default_factory=lambda: os.getenv("VOLCENGINE_ACCESS_KEY")
52+
)
53+
volcengine_secret_key: str | None = Field(
54+
default_factory=lambda: os.getenv("VOLCENGINE_SECRET_KEY")
55+
)
56+
tos_vector_bucket_name: str | None = Field(
57+
default_factory=lambda: os.getenv("DATABASE_TOS_VECTOR_BUCKET")
58+
)
59+
tos_vector_account_id: str | None = Field(
60+
default_factory=lambda: os.getenv("DATABASE_TOS_VECTOR_ACCOUNT_ID")
61+
)
62+
tos_vector_config: TOSVectorConfig = Field(default_factory=TOSVectorConfig)
63+
64+
session_token: str = ""
65+
66+
embedding_config: EmbeddingModelConfig | NormalEmbeddingModelConfig = Field(
67+
default_factory=EmbeddingModelConfig
68+
)
69+
70+
def model_post_init(self, __context: Any) -> None:
71+
self.precheck_index_naming()
72+
self._tos_vector_client = VectorClient(
73+
ak=self.volcengine_access_key,
74+
sk=self.volcengine_secret_key,
75+
**self.tos_vector_config.model_dump(),
76+
)
77+
# create_bucket and index if not exist
78+
self._create_index()
79+
80+
self._embed_model = OpenAILikeEmbedding(
81+
model_name=self.embedding_config.name,
82+
api_key=self.embedding_config.api_key,
83+
api_base=self.embedding_config.api_base,
84+
)
85+
86+
def _index_exists(self) -> bool:
87+
try:
88+
index_exist = self._tos_vector_client.get_index(
89+
vector_bucket_name=self.tos_vector_bucket_name,
90+
account_id=self.tos_vector_account_id,
91+
index_name=self.index,
92+
)
93+
return index_exist.status_code == 200
94+
except TosServerError as e:
95+
if e.status_code == 404:
96+
return False
97+
else:
98+
raise e
99+
100+
def _split_documents(self, documents: list[Document]) -> list[BaseNode]:
101+
"""Split document into chunks"""
102+
nodes = []
103+
for document in documents:
104+
splitter = get_llama_index_splitter(document.metadata.get("file_path", ""))
105+
_nodes = splitter.get_nodes_from_documents([document])
106+
nodes.extend(_nodes)
107+
return nodes
108+
109+
def _create_index(self):
110+
# no need to check if bucket exists, create_bucket will create it if not exist
111+
self._tos_vector_client.create_vector_bucket(
112+
vector_bucket_name=self.tos_vector_bucket_name,
113+
)
114+
115+
if not self._index_exists():
116+
self._tos_vector_client.create_index(
117+
vector_bucket_name=self.tos_vector_bucket_name,
118+
account_id=self.tos_vector_account_id,
119+
index_name=self.index,
120+
data_type=DataType.DataTypeFloat32,
121+
dimension=self.embedding_config.dim,
122+
distance_metric=DistanceMetricType.DistanceMetricCosine,
123+
)
124+
125+
def precheck_index_naming(self) -> None:
126+
pass
127+
128+
def _process_and_store_documents(self, documents: list[Document]) -> bool:
129+
nodes = self._split_documents(documents)
130+
vectors = []
131+
for node in nodes:
132+
if not node.text:
133+
continue
134+
embedding = self._embed_model.get_text_embedding(node.text)
135+
vectors.append(
136+
Vector(
137+
key=node.node_id,
138+
data=VectorData(float32=embedding),
139+
metadata={"text": node.text, "metadata": json.dumps(node.metadata)},
140+
)
141+
)
142+
result = self._tos_vector_client.put_vectors(
143+
vector_bucket_name=self.tos_vector_bucket_name,
144+
account_id=self.tos_vector_account_id,
145+
index_name=self.index,
146+
vectors=vectors,
147+
)
148+
return result.status_code == 200
149+
150+
@override
151+
def add_from_directory(self, directory: str, *args, **kwargs) -> bool:
152+
# fixme
153+
logger.warning(
154+
"add_from_directory is not yet fully developed and may have issues such as missing images."
155+
)
156+
documents = SimpleDirectoryReader(input_dir=directory).load_data()
157+
return self._process_and_store_documents(documents)
158+
159+
@override
160+
def add_from_files(self, files: list[str], *args, **kwargs) -> bool:
161+
# fixme
162+
logger.warning(
163+
"add_from_files is not yet fully developed and may have issues such as missing images."
164+
)
165+
documents = SimpleDirectoryReader(input_files=files).load_data()
166+
return self._process_and_store_documents(documents)
167+
168+
@override
169+
def add_from_text(self, text: str | list[str], *args, **kwargs) -> bool:
170+
if isinstance(text, str):
171+
documents = [Document(text=text)]
172+
else:
173+
documents = [Document(text=t) for t in text]
174+
175+
return self._process_and_store_documents(documents)
176+
177+
@override
178+
def search(self, query: str, top_k: int = 5) -> list[str]:
179+
query_vector = self._embed_model.get_text_embedding(query)
180+
181+
search_result = self._tos_vector_client.query_vectors(
182+
vector_bucket_name=self.tos_vector_bucket_name,
183+
account_id=self.tos_vector_account_id,
184+
index_name=self.index,
185+
query_vector=VectorData(float32=query_vector),
186+
top_k=top_k,
187+
return_metadata=True,
188+
)
189+
190+
return [vector.metadata["text"] for vector in search_result.vectors]

‎veadk/knowledgebase/knowledgebase.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,12 @@ def _get_backend_cls(backend: str) -> type[BaseKnowledgebaseBackend]:
5151
)
5252

5353
return RedisKnowledgeBackend
54+
case "tos_vector":
55+
from veadk.knowledgebase.backends.tos_vector_backend import (
56+
TosVectorKnowledgeBackend,
57+
)
58+
59+
return TosVectorKnowledgeBackend
5460

5561
raise ValueError(f"Unsupported knowledgebase backend: {backend}")
5662

@@ -165,7 +171,8 @@ class KnowledgeBase(BaseModel):
165171
description: str = "This knowledgebase stores some user-related information."
166172

167173
backend: Union[
168-
Literal["local", "opensearch", "viking", "redis"], BaseKnowledgebaseBackend
174+
Literal["local", "opensearch", "viking", "redis", "tos_vector"],
175+
BaseKnowledgebaseBackend,
169176
] = "local"
170177

171178
backend_config: dict = Field(default_factory=dict)

0 commit comments

Comments
 (0)