484 lines
17 KiB
Python
484 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
百度VDB向量数据库后端
|
||
支持多模态向量存储和检索
|
||
"""
|
||
|
||
import pymochow
|
||
from pymochow.configuration import Configuration
|
||
from pymochow.auth.bce_credentials import BceCredentials
|
||
from pymochow.model.schema import Schema, Field, VectorIndex, HNSWParams, SecondaryIndex
|
||
from pymochow.model.enum import FieldType, IndexType, MetricType
|
||
from pymochow.model.table import Partition, Row, VectorTopkSearchRequest, FloatVector, VectorSearchConfig
|
||
import numpy as np
|
||
import json
|
||
import time
|
||
import logging
|
||
from typing import List, Tuple, Union, Dict, Any
|
||
import hashlib
|
||
import os
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class BaiduVDBBackend:
|
||
"""百度VDB向量数据库后端"""
|
||
|
||
def __init__(self, account: str = "root", api_key: str = "vdb$yjr9ln3n0td",
|
||
endpoint: str = "http://180.76.96.191:5287", database_name: str = "multimodal_retrieval"):
|
||
"""
|
||
初始化百度VDB后端
|
||
|
||
Args:
|
||
account: 用户名
|
||
api_key: API密钥
|
||
endpoint: 服务器端点
|
||
database_name: 数据库名称
|
||
"""
|
||
self.account = account
|
||
self.api_key = api_key
|
||
self.endpoint = endpoint
|
||
self.database_name = database_name
|
||
|
||
# 初始化客户端
|
||
self.client = None
|
||
self.db = None
|
||
self.text_table = None
|
||
self.image_table = None
|
||
|
||
# 表名配置
|
||
self.text_table_name = "text_vectors"
|
||
self.image_table_name = "image_vectors"
|
||
|
||
# 向量维度（根据Ops-MM-embedding-v1-7B模型）
|
||
self.vector_dimension = 3584
|
||
|
||
self._connect()
|
||
self._ensure_database()
|
||
self._ensure_tables()
|
||
|
||
def _connect(self):
|
||
"""连接到百度VDB"""
|
||
try:
|
||
config = Configuration(
|
||
credentials=BceCredentials(self.account, self.api_key),
|
||
endpoint=self.endpoint
|
||
)
|
||
self.client = pymochow.MochowClient(config)
|
||
logger.info(f"✅ 成功连接到百度VDB: {self.endpoint}")
|
||
except Exception as e:
|
||
logger.error(f"❌ 连接百度VDB失败: {e}")
|
||
raise
|
||
|
||
def _ensure_database(self):
|
||
"""确保数据库存在"""
|
||
try:
|
||
# 检查数据库是否存在
|
||
db_list = self.client.list_databases()
|
||
existing_dbs = [db.database_name for db in db_list]
|
||
|
||
if self.database_name not in existing_dbs:
|
||
logger.info(f"创建数据库: {self.database_name}")
|
||
self.db = self.client.create_database(self.database_name)
|
||
else:
|
||
logger.info(f"使用现有数据库: {self.database_name}")
|
||
self.db = self.client.database(self.database_name)
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 数据库操作失败: {e}")
|
||
raise
|
||
|
||
def _ensure_tables(self):
|
||
"""确保表存在"""
|
||
try:
|
||
# 获取现有表列表
|
||
existing_tables = self.db.list_table()
|
||
existing_table_names = [table.table_name for table in existing_tables]
|
||
|
||
# 创建文本向量表
|
||
if self.text_table_name not in existing_table_names:
|
||
self._create_text_table()
|
||
else:
|
||
self.text_table = self.db.table(self.text_table_name)
|
||
logger.info(f"使用现有文本表: {self.text_table_name}")
|
||
|
||
# 创建图像向量表
|
||
if self.image_table_name not in existing_table_names:
|
||
self._create_image_table()
|
||
else:
|
||
self.image_table = self.db.table(self.image_table_name)
|
||
logger.info(f"使用现有图像表: {self.image_table_name}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 表操作失败: {e}")
|
||
raise
|
||
|
||
def _create_text_table(self):
|
||
"""创建文本向量表"""
|
||
try:
|
||
logger.info(f"创建文本向量表: {self.text_table_name}")
|
||
|
||
# 定义字段 - 移除可能导致问题的复杂配置
|
||
fields = [
|
||
Field("id", FieldType.STRING, primary_key=True, not_null=True),
|
||
Field("text_content", FieldType.STRING, not_null=True),
|
||
Field("vector", FieldType.FLOAT_VECTOR, not_null=True, dimension=self.vector_dimension)
|
||
]
|
||
|
||
# 定义索引 - 简化配置
|
||
indexes = [
|
||
VectorIndex(
|
||
index_name="text_vector_idx",
|
||
index_type=IndexType.HNSW,
|
||
field="vector",
|
||
metric_type=MetricType.COSINE,
|
||
params=HNSWParams(m=16, efconstruction=100),
|
||
auto_build=True
|
||
)
|
||
]
|
||
|
||
# 创建表 - 简化配置
|
||
self.text_table = self.db.create_table(
|
||
table_name=self.text_table_name,
|
||
replication=1, # 单副本
|
||
schema=Schema(fields=fields, indexes=indexes)
|
||
)
|
||
|
||
logger.info(f"✅ 文本向量表创建成功")
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 创建文本表失败: {e}")
|
||
raise
|
||
|
||
def _create_image_table(self):
|
||
"""创建图像向量表"""
|
||
try:
|
||
logger.info(f"创建图像向量表: {self.image_table_name}")
|
||
|
||
# 定义字段 - 移除可能导致问题的复杂配置
|
||
fields = [
|
||
Field("id", FieldType.STRING, primary_key=True, not_null=True),
|
||
Field("image_path", FieldType.STRING, not_null=True),
|
||
Field("vector", FieldType.FLOAT_VECTOR, not_null=True, dimension=self.vector_dimension)
|
||
]
|
||
|
||
# 定义索引 - 简化配置
|
||
indexes = [
|
||
VectorIndex(
|
||
index_name="image_vector_idx",
|
||
index_type=IndexType.HNSW,
|
||
field="vector",
|
||
metric_type=MetricType.COSINE,
|
||
params=HNSWParams(m=16, efconstruction=100),
|
||
auto_build=True
|
||
)
|
||
]
|
||
|
||
# 创建表 - 简化配置
|
||
self.image_table = self.db.create_table(
|
||
table_name=self.image_table_name,
|
||
replication=1, # 单副本
|
||
schema=Schema(fields=fields, indexes=indexes)
|
||
)
|
||
|
||
logger.info(f"✅ 图像向量表创建成功")
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 创建图像表失败: {e}")
|
||
raise
|
||
|
||
def _generate_id(self, content: str) -> str:
|
||
"""生成唯一ID"""
|
||
return hashlib.md5(f"{content}_{time.time()}".encode()).hexdigest()
|
||
|
||
def store_text_vectors(self, texts: List[str], vectors: np.ndarray, metadata: List[Dict] = None) -> List[str]:
|
||
"""
|
||
存储文本向量
|
||
|
||
Args:
|
||
texts: 文本列表
|
||
vectors: 向量数组
|
||
metadata: 元数据列表
|
||
|
||
Returns:
|
||
存储的ID列表
|
||
"""
|
||
if len(texts) != len(vectors):
|
||
raise ValueError("文本数量与向量数量不匹配")
|
||
|
||
try:
|
||
rows = []
|
||
ids = []
|
||
current_time = int(time.time() * 1000) # 毫秒时间戳
|
||
|
||
for i, (text, vector) in enumerate(zip(texts, vectors)):
|
||
doc_id = self._generate_id(text)
|
||
ids.append(doc_id)
|
||
|
||
# 准备元数据
|
||
meta = metadata[i] if metadata and i < len(metadata) else {}
|
||
meta_json = json.dumps(meta, ensure_ascii=False)
|
||
|
||
row = Row(
|
||
id=doc_id,
|
||
text_content=text,
|
||
vector=vector.tolist()
|
||
)
|
||
rows.append(row)
|
||
|
||
# 批量插入
|
||
self.text_table.upsert(rows)
|
||
logger.info(f"✅ 成功存储 {len(texts)} 条文本向量")
|
||
|
||
return ids
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 存储文本向量失败: {e}")
|
||
raise
|
||
|
||
def store_image_vectors(self, image_paths: List[str], vectors: np.ndarray, metadata: List[Dict] = None) -> List[str]:
|
||
"""
|
||
存储图像向量
|
||
|
||
Args:
|
||
image_paths: 图像路径列表
|
||
vectors: 向量数组
|
||
metadata: 元数据列表
|
||
|
||
Returns:
|
||
存储的ID列表
|
||
"""
|
||
if len(image_paths) != len(vectors):
|
||
raise ValueError("图像数量与向量数量不匹配")
|
||
|
||
try:
|
||
rows = []
|
||
ids = []
|
||
current_time = int(time.time() * 1000) # 毫秒时间戳
|
||
|
||
for i, (image_path, vector) in enumerate(zip(image_paths, vectors)):
|
||
doc_id = self._generate_id(image_path)
|
||
ids.append(doc_id)
|
||
|
||
# 准备元数据
|
||
meta = metadata[i] if metadata and i < len(metadata) else {}
|
||
meta_json = json.dumps(meta, ensure_ascii=False)
|
||
|
||
row = Row(
|
||
id=doc_id,
|
||
image_path=image_path,
|
||
vector=vector.tolist()
|
||
)
|
||
rows.append(row)
|
||
|
||
# 批量插入
|
||
self.image_table.upsert(rows)
|
||
logger.info(f"✅ 成功存储 {len(image_paths)} 条图像向量")
|
||
|
||
return ids
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 存储图像向量失败: {e}")
|
||
raise
|
||
|
||
def search_text_vectors(self, query_vector: np.ndarray, top_k: int = 5,
|
||
filter_condition: str = None) -> List[Tuple[str, str, float, Dict]]:
|
||
"""
|
||
搜索文本向量
|
||
|
||
Args:
|
||
query_vector: 查询向量
|
||
top_k: 返回结果数量
|
||
filter_condition: 过滤条件
|
||
|
||
Returns:
|
||
(id, text_content, score, metadata) 列表
|
||
"""
|
||
try:
|
||
# 构建搜索请求
|
||
request = VectorTopkSearchRequest(
|
||
vector_field="vector",
|
||
vector=FloatVector(query_vector.tolist()),
|
||
limit=top_k,
|
||
filter=filter_condition,
|
||
config=VectorSearchConfig(ef=200)
|
||
)
|
||
|
||
# 执行搜索
|
||
results = self.text_table.vector_search(request=request)
|
||
|
||
# 解析结果
|
||
search_results = []
|
||
for result in results:
|
||
doc_id = result.get('id', '')
|
||
text_content = result.get('text_content', '')
|
||
score = result.get('_score', 0.0)
|
||
|
||
search_results.append((doc_id, text_content, float(score), {}))
|
||
|
||
logger.info(f"✅ 文本向量搜索完成，返回 {len(search_results)} 条结果")
|
||
return search_results
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 文本向量搜索失败: {e}")
|
||
return []
|
||
|
||
def search_image_vectors(self, query_vector: np.ndarray, top_k: int = 5,
|
||
filter_condition: str = None) -> List[Tuple[str, str, str, float, Dict]]:
|
||
"""
|
||
搜索图像向量
|
||
|
||
Args:
|
||
query_vector: 查询向量
|
||
top_k: 返回结果数量
|
||
filter_condition: 过滤条件
|
||
|
||
Returns:
|
||
(id, image_path, image_name, score, metadata) 列表
|
||
"""
|
||
try:
|
||
# 构建搜索请求
|
||
request = VectorTopkSearchRequest(
|
||
vector_field="vector",
|
||
vector=FloatVector(query_vector.tolist()),
|
||
limit=top_k,
|
||
filter=filter_condition,
|
||
config=VectorSearchConfig(ef=200)
|
||
)
|
||
|
||
# 执行搜索
|
||
results = self.image_table.vector_search(request=request)
|
||
|
||
# 解析结果
|
||
search_results = []
|
||
for result in results:
|
||
doc_id = result.get('id', '')
|
||
image_path = result.get('image_path', '')
|
||
image_name = os.path.basename(image_path)
|
||
score = result.get('_score', 0.0)
|
||
|
||
search_results.append((doc_id, image_path, image_name, float(score), {}))
|
||
|
||
logger.info(f"✅ 图像向量搜索完成，返回 {len(search_results)} 条结果")
|
||
return search_results
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 图像向量搜索失败: {e}")
|
||
return []
|
||
|
||
def get_statistics(self) -> Dict[str, Any]:
|
||
"""获取数据库统计信息"""
|
||
try:
|
||
stats = {}
|
||
|
||
# 文本表统计
|
||
if self.text_table:
|
||
text_stats = self.text_table.stats()
|
||
stats['text_table'] = {
|
||
'row_count': text_stats.get('rowCount', 0),
|
||
'memory_size_mb': text_stats.get('memorySizeInByte', 0) / (1024 * 1024),
|
||
'disk_size_mb': text_stats.get('diskSizeInByte', 0) / (1024 * 1024)
|
||
}
|
||
|
||
# 图像表统计
|
||
if self.image_table:
|
||
image_stats = self.image_table.stats()
|
||
stats['image_table'] = {
|
||
'row_count': image_stats.get('rowCount', 0),
|
||
'memory_size_mb': image_stats.get('memorySizeInByte', 0) / (1024 * 1024),
|
||
'disk_size_mb': image_stats.get('diskSizeInByte', 0) / (1024 * 1024)
|
||
}
|
||
|
||
return stats
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 获取统计信息失败: {e}")
|
||
return {}
|
||
|
||
def clear_all_data(self):
|
||
"""清空所有数据"""
|
||
try:
|
||
# 清空文本表
|
||
if self.text_table:
|
||
self.text_table.delete(filter="*")
|
||
logger.info("✅ 文本表数据已清空")
|
||
|
||
# 清空图像表
|
||
if self.image_table:
|
||
self.image_table.delete(filter="*")
|
||
logger.info("✅ 图像表数据已清空")
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 清空数据失败: {e}")
|
||
raise
|
||
|
||
def close(self):
|
||
"""关闭连接"""
|
||
try:
|
||
if self.client:
|
||
self.client.close()
|
||
logger.info("✅ 百度VDB连接已关闭")
|
||
except Exception as e:
|
||
logger.error(f"❌ 关闭连接失败: {e}")
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
self.close()
|
||
|
||
|
||
def test_vdb_backend():
|
||
"""测试VDB后端功能"""
|
||
print("=" * 60)
|
||
print("测试百度VDB后端功能")
|
||
print("=" * 60)
|
||
|
||
try:
|
||
# 初始化后端
|
||
vdb = BaiduVDBBackend()
|
||
|
||
# 测试数据
|
||
test_texts = [
|
||
"这是一个测试文本",
|
||
"多模态检索系统",
|
||
"向量数据库测试"
|
||
]
|
||
|
||
# 生成测试向量（随机向量）
|
||
test_vectors = np.random.rand(3, 3584).astype(np.float32)
|
||
|
||
# 存储文本向量
|
||
print("1. 存储文本向量...")
|
||
text_ids = vdb.store_text_vectors(test_texts, test_vectors)
|
||
print(f" 存储成功，ID: {text_ids}")
|
||
|
||
# 搜索文本向量
|
||
print("\n2. 搜索文本向量...")
|
||
query_vector = np.random.rand(3584).astype(np.float32)
|
||
results = vdb.search_text_vectors(query_vector, top_k=3)
|
||
print(f" 搜索结果: {len(results)} 条")
|
||
for i, (doc_id, text, score, meta) in enumerate(results, 1):
|
||
print(f" {i}. {text[:30]}... (相似度: {score:.4f})")
|
||
|
||
# 获取统计信息
|
||
print("\n3. 数据库统计信息...")
|
||
stats = vdb.get_statistics()
|
||
print(f" 统计信息: {stats}")
|
||
|
||
# 清理测试数据
|
||
print("\n4. 清理测试数据...")
|
||
vdb.clear_all_data()
|
||
print(" 清理完成")
|
||
|
||
print("\n✅ 百度VDB后端测试完成！")
|
||
|
||
except Exception as e:
|
||
print(f"\n❌ 测试失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
test_vdb_backend()
|