# -*- coding:utf-8 -*-
# 导入 BGE-M3 嵌入函数，用于生成文档和查询的向量表示
import torch.cuda
from langchain_core.documents import Document
from milvus_model.hybrid import BGEM3EmbeddingFunction
# 导入 Milvus 相关类，用于操作向量数据库
from pymilvus import MilvusClient, DataType, AnnSearchRequest, WeightedRanker
# 导入 Document 类，用于创建文档对象
# from langchain.docstore.document import Document
# 导入 CrossEncoder，用于重排序和 NLI 判断
from sentence_transformers import CrossEncoder
# 导入 hashlib 模块，用于生成唯一 ID 的哈希值
import hashlib
from rag_qa.core.document_processor import *
import sys, os
# 获取当前文件所在目录的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))
# print(f'current_dir--》{current_dir}')
# 获取core文件所在的目录的绝对路径
rag_qa_path = os.path.dirname(current_dir)
# print(f'rag_qa_path--》{rag_qa_path}')
sys.path.insert(0, rag_qa_path)
# 获取根目录文件所在的绝对位置
project_root = os.path.dirname(rag_qa_path)
sys.path.insert(0, project_root)
from base import logger, Config
conf = Config()
# core/vector_store.py
# 定义 VectorStore 类，封装向量存储和检索功能
class VectorStore:
# 初始化方法，设置向量存储的基本参数
def __init__(self,
collection_name=conf.MILVUS_COLLECTION_NAME,
host=conf.MILVUS_HOST,
port=conf.MILVUS_PORT,
database=conf.MILVUS_DATABASE_NAME):
# 设置 Milvus 集合名称
self.collection_name = collection_name
# 设置 Milvus 主机地址
self.host = host
# 设置 Milvus 端口号
self.port = port
# 设置 Milvus 数据库名称
self.database = database
# 设置日志记录器
self.logger = logger
# 检查CUDA是否可用
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
# 日志提醒使用的是什么设备
self.logger.info(f"使用设置：{self.device}")
# 初始化 BGE-Reranker 模型，用于重排序检索结果
reranker_path = os.path.join(rag_qa_path, 'models', 'bge-reranker-large')
# print(f'reranker_path--》{reranker_path}')
self.reranker = CrossEncoder(reranker_path, device=self.device)
# 初始化 BGE-M3 嵌入函数，使用 CPU 设备，不启用 FP16
m3_path = os.path.join(rag_qa_path, 'models', 'bge-m3')
self.embedding_function = BGEM3EmbeddingFunction(model_name_or_path=m3_path, use_fp16=(self.device == 'cuda'),
device=self.device)
# 获取稠密向量的维度# 1024
self.dense_dim = self.embedding_function.dim["dense"]
# 初始化 Milvus 客户端，连接到指定主机和数据库
self.client = MilvusClient(uri=f"http://{self.host}:{self.port}", db_name=self.database)
# 调用方法创建或加载 Milvus 集合
self._create_or_load_collection()
# 类私有化方法
def _create_or_load_collection(self):
# 检查指定集合是否已经存在
if not self.client.has_collection(self.collection_name):
# 创建集合 Schema，禁用自动 ID，启用动态字段
schema = self.client.create_schema(auto_id=False, enable_dynamic_field=True)
# 添加 ID 字段，作为主键，VARCHAR 类型，最大长度 100
schema.add_field(field_name="id", datatype=DataType.VARCHAR, is_primary=True, max_length=100)
# 添加文本字段，VARCHAR 类型，最大长度 65535
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=65535)
# 添加稠密向量字段，FLOAT_VECTOR 类型，维度由嵌入函数指定
schema.add_field(field_name="dense_vector", datatype=DataType.FLOAT_VECTOR, dim=self.dense_dim)
# 添加稀疏向量字段，SPARSE_FLOAT_VECTOR 类型
schema.add_field(field_name="sparse_vector", datatype=DataType.SPARSE_FLOAT_VECTOR)
# 添加父块 ID 字段，VARCHAR 类型，最大长度 100
schema.add_field(field_name="parent_id", datatype=DataType.VARCHAR, max_length=100)
# 添加父块内容字段，VARCHAR 类型，最大长度 65535
schema.add_field(field_name="parent_content", datatype=DataType.VARCHAR, max_length=65535)
# 添加学科类别字段，VARCHAR 类型，最大长度 50
schema.add_field(field_name="source", datatype=DataType.VARCHAR, max_length=50)
# 添加时间戳字段，VARCHAR 类型，最大长度 50
schema.add_field(field_name="timestamp", datatype=DataType.VARCHAR, max_length=50)
# 创建索引参数对象
index_params = self.client.prepare_index_params()
# 为稠密向量字段添加 IVF_FLAT 索引，度量类型为内积 (IP)
index_params.add_index(
field_name="dense_vector",
index_name="dense_index",
index_type="IVF_FLAT",
metric_type="IP",
params={"nlist": 128}
)
# 为稀疏向量字段添加 SPARSE_INVERTED_INDEX 索引，度量类型为内积 (IP)
index_params.add_index(
field_name="sparse_vector",
index_name="sparse_index",
index_type="SPARSE_INVERTED_INDEX",
metric_type="IP",
params={"drop_ratio_build": 0.2}
)
# 创建 Milvus 集合，应用定义的 Schema 和索引参数
self.client.create_collection(collection_name=self.collection_name, schema=schema,
index_params=index_params)
# 记录创建集合的日志
logger.info(f"已创建集合 {self.collection_name}")
# 如果集合已存在
else:
# 记录加载集合的日志
logger.info(f"已加载集合 {self.collection_name}")
# 将集合加载到内存，确保可立即查询
self.client.load_collection(self.collection_name)
# 定义方法，往向量存储添加文档
def add_documents(self, documents):
# print(f'documents--》{documents[0]}')
# 提取所有文档的内容列表
texts = [doc.page_content for doc in documents]
# 使用 BGE-M3 嵌入函数生成文档的嵌入
embeddings = self.embedding_function(texts)
# print(f'embeddings-->', embeddings)
# print(f'embeddings.dense-->', embeddings['dense'])
# print(f'embeddings.sparse-->', embeddings['sparse'])
# 初始化空列表，存储插入的数据
data = []
# 遍历每个文档，带上索引i
for i, doc in enumerate(documents):
# 生成文档内容的哈希值作为唯一的ID
text_hash = hashlib.md5(doc.page_content.encode('utf-8')).hexdigest()
# print(f'text_hash--》{text_hash}')
# print(f'text_hash--》{type(text_hash)}')
# 初始化一个稀疏向量的字典（Milvus要求存储稀疏向量的格式）
sparse_vector = {}
try:
# 新版本 milvus-model 使用 coo_array 格式
row = embeddings["sparse"][i]
print('i-->',i, "row-->", row)
if hasattr(row, 'col'): # coo_array 格式
indices = row.col
values = row.data
else: # csr_matrix 格式
indices = row.indices
values = row.data
except Exception as e:
# 兼容旧版本 milvus-model
row = embeddings["sparse"].getrow(0)
indices = row.indices
values = row.data
# 获取第i行对应的稀疏向量数据[0.4, 0.2, 0, 0, 0.1]
# row = embeddings["sparse"].getrow(i)
# row = embeddings["sparse"][i] # 新版本milvus-model，支持这种获取稀疏向量的形式
# indices = row.row
# print(f'row--》{row}')
# print(f'row--》{row.shape}')
# 获取稀疏向量的非零值的索引
# indices = row.col
# print(f'indics--》{indics}')
# 获取稀疏向量的非零值
# values = row.data
# 将索引和值进行配对，存储到字典中
for idx, value in zip(indices, values):
sparse_vector[idx] = value
# print(f'sparse_vector--》{sparse_vector}')
# print(f'sparse_vector--》{len(sparse_vector)}')
# print(embeddings["dense"][i])
# print(embeddings["dense"][i].shape)
# 创建数据字典，包含所有字段
data.append({
"id": text_hash,
"text": doc.page_content,
"dense_vector": embeddings["dense"][i],
"sparse_vector": sparse_vector,
"parent_id": doc.metadata["parent_id"],
"parent_content": doc.metadata["parent_content"],
"source": doc.metadata.get("source", "unknown"),
"timestamp": doc.metadata.get("timestamp", "unknown")
})
# print("data-->", data)
# 检查是否有数据需要插入
if data:
# 使用 upsert 操作插入数据，覆盖重复 ID
self.client.upsert(collection_name=self.collection_name, data=data)
# 记录插入或更新的文档数量日志
logger.info(f"已插入或更新 {len(data)} 个文档")
# 定义方法，执行混合检索并重排序
def hybrid_search_with_rerank(self, query, k=conf.RETRIEVAL_K, source_filter=None):
# 使用 BGE-M3 嵌入函数生成查询的嵌入
query_embeddings = self.embedding_function([query])
# 获取查询的稠密向量
dense_query_vector = query_embeddings["dense"][0]
# print(f'dense_query_vector--》{dense_query_vector.shape}')
# 初始化查询的稀疏向量字典
sparse_query_vector = {}
# ====================
# 获取查询稀疏向量的第 0 行数据
# row = query_embeddings["sparse"][0]
# # 获取稀疏向量的非零值索引
# indices = row.indices
# # 获取稀疏向量的非零值
# values = row.data
# ====================
try:
# 新版本 milvus-model 使用 coo_array 格式
row = query_embeddings["sparse"][0]
if hasattr(row, 'col'): # coo_array 格式
indices = row.col
values = row.data
else: # csr_matrix 格式
indices = row.indices
values = row.data
except Exception as e:
# 兼容旧版本 milvus-model
row = query_embeddings["sparse"].getrow(0)
indices = row.indices
values = row.data
# 将索引和值配对，填充稀疏向量字典
for idx, value in zip(indices, values):
sparse_query_vector[idx] = value
# print(f'sparse_query_vector-->{sparse_query_vector}')
# 初始化过滤表达式，默认不过滤
filter_expr = f"source == '{source_filter}'" if source_filter else ""
# print(f'filter_expr--》{filter_expr}')
# 创建稠密向量搜索请求
dense_request = AnnSearchRequest(
data=[dense_query_vector],
anns_field="dense_vector",
param={"metric_type": "IP", "params": {"nprobe": 10}},
limit=k,
expr=filter_expr
)
# 创建稀疏向量搜索请求
sparse_request = AnnSearchRequest(
data=[sparse_query_vector],
anns_field="sparse_vector",
param={"metric_type": "IP", "params": {}},
limit=k,
expr=filter_expr
)
# 创建加权排序器，稠密向量权重 1.0，稀疏向量权重 0.7
ranker = WeightedRanker(1.0, 0.7)
# 执行混合搜索，返回 Top-K 结果
results = self.client.hybrid_search(
collection_name=self.collection_name,
reqs=[dense_request, sparse_request],
ranker=ranker,
limit=k,
output_fields=["text", "parent_id", "parent_content", "source", "timestamp"]
)[0]
# print(f'results--》{results}')
# print(f'results--》{type(results)}')
# print(f'results--》{len(results)}')
# 将上述搜索到的结果进行Document对象封装，便于查询使用
sub_chunks = [self._doc_from_hit(hit["entity"]) for hit in results]
# print(f'sub_chunks--》{len(sub_chunks)}')
# 从子块中提取去重的父文档
parent_docs = self._get_unique_parent_docs(sub_chunks)
# print(f'parent_docs--》{parent_docs}')
# print(f'parent_docs--》{len(parent_docs)}')
# # 如果只有1个文档或者没有，直接返回跳过重排序
if len(parent_docs) < 2:
return parent_docs[:conf.CANDIDATE_M]
# 如果有父文档，进行重排序
if parent_docs:
# 创建查询与文档内容的配对列表
pairs = [[query, doc.page_content] for doc in parent_docs]
# 使用 BGE-Reranker 计算每个配对的得分
scores = self.reranker.predict(pairs)
# print(f'scores--》{scores}')
# 根据得分从高到低排序文档
ranked_parent_docs = [doc for _, doc in sorted(zip(scores, parent_docs), reverse=True)]
# 如果没有父文档，返回空列表
else:
ranked_parent_docs = []
# 返回前 m 个重排序后的文档
return ranked_parent_docs[:conf.CANDIDATE_M]
def _get_unique_parent_docs(self, sub_chunks):
# 初始化集合，用于存储已处理的父块内容（去重）
parent_contents = set()
# 初始化列表，用于存储唯一父文档
unique_docs = []
# 遍历所有子块
for chunk in sub_chunks:
# 获取子块的父块内容，默认为子块内容
parent_content = chunk.metadata.get("parent_content", chunk.page_content)
# 检查父块内容是否非空且未重复
if parent_content and parent_content not in parent_contents:
# 创建新的 Document 对象，包含父块内容和元数据
unique_docs.append(Document(page_content=parent_content, metadata=chunk.metadata))
# 将父块内容添加到去重集合
parent_contents.add(parent_content)
# 返回去重后的父文档列表
return unique_docs
# 定义类似私有方法，从 Milvus 查询结果创建 Document 对象
def _doc_from_hit(self, hit):
# 创建并返回 Document 对象，填充内容和元数据
return Document(
page_content=hit.get("text"),
metadata={
"parent_id": hit.get("parent_id"),
"parent_content": hit.get("parent_content"),
"source": hit.get("source"),
"timestamp": hit.get("timestamp")
}
)
def test_bge_m3():
# 准备query
query = "课程优势是什么"
# 测试BGE-M3模型
bgem3 = BGEM3EmbeddingFunction(model_name_or_path=r"D:\code\workspace7\integrated_qa_system\rag_qa\models\bge-m3",
use_fp16=False,
device="cpu")
# 把query送给模型，得到向量
result = bgem3([query])
# 打印结果
print("result-->", result)
if __name__ == "__main__":
# TODO 需要先在milvus中创建itcast数据库，否则执行代码会报错。
vector_store = VectorStore()
# 把知识库文件加载到Milvus向量数据库中
# directory_path = 'D:/code/workspace7/integrated_qa_system/rag_qa/data/ai_data'
# print(f"embedding_function.dim--》{vector_store.embedding_function.dim}")
# documents = process_documents(directory_path)
# print("documents-->", documents)
# vector_store.add_documents(documents)
# 测试代码
query = "课程新增了什么技术"
results = vector_store.hybrid_search_with_rerank(query, source_filter='ai')
print(f'results-->{results}')
print(f'results-->{len(results)}')