147 lines
4.9 KiB
Python
147 lines
4.9 KiB
Python
from datetime import datetime
|
||
from typing import Optional, List, Dict, Any
|
||
|
||
from pymilvus import MilvusClient, DataType
|
||
|
||
from app.client.milvus_client_manager import milvus_client
|
||
from app.conf.app_config import app_config
|
||
from app.core.log import logger
|
||
|
||
|
||
class QARepository:
|
||
collection_name: str = 'message_qa'
|
||
|
||
def __init__(self):
|
||
self.dim = app_config.milvus.embedding_size
|
||
self._ensure_connection()
|
||
|
||
def _ensure_connection(self):
|
||
if milvus_client.client is None:
|
||
milvus_client.init()
|
||
|
||
def create_collection(self, dim: int = None) -> bool:
|
||
"""
|
||
创建消息摘要集合
|
||
|
||
Args:
|
||
dim: 向量维度，默认从配置读取
|
||
|
||
Returns:
|
||
是否成功创建（已存在返回False）
|
||
"""
|
||
if dim is not None:
|
||
self.dim = dim
|
||
|
||
if milvus_client.has_collection(self.collection_name):
|
||
return False
|
||
|
||
schema = milvus_client.client.create_schema(
|
||
auto_id=True,
|
||
enable_dynamic_field=True,
|
||
description="QA向量集合"
|
||
)
|
||
|
||
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(field_name="question", datatype=DataType.VARCHAR, max_length=1024)
|
||
schema.add_field(field_name="answer", datatype=DataType.VARCHAR, max_length=65535)
|
||
schema.add_field(field_name="parent_id", datatype=DataType.INT64)
|
||
schema.add_field(field_name="question_dense_vector", datatype=DataType.FLOAT_VECTOR, dim=1024)
|
||
schema.add_field(field_name="created_at", datatype=DataType.INT64)
|
||
|
||
index_params = milvus_client.client.prepare_index_params()
|
||
index_params.add_index(
|
||
field_name="question_dense_vector",
|
||
index_type="HNSW",
|
||
metric_type="COSINE",
|
||
params={"M": 32, "efConstruction": 300}
|
||
)
|
||
milvus_client.client.create_collection(
|
||
collection_name=self.collection_name,
|
||
schema=schema,
|
||
index_params=index_params
|
||
)
|
||
|
||
return True
|
||
|
||
async def save_qa_to_milvus(self, qa_list: List[Dict[str, Any]]) -> List[int]:
|
||
"""
|
||
将QA数据入库到Milvus
|
||
|
||
Args:
|
||
qa_pairs: QA对列表，格式 [{"question": "...", "answer": "..."}, ...]
|
||
parent_id: 父记录ID，默认为0
|
||
|
||
Returns:
|
||
插入的记录ID列表
|
||
"""
|
||
try:
|
||
# 构建入库数据
|
||
created_at = int(datetime.now().timestamp())
|
||
data = []
|
||
for qa in qa_list:
|
||
data.append({
|
||
"question": qa.get("question", ""),
|
||
"answer": qa.get("answer", ""),
|
||
"parent_id": qa.get("parent_id", 0),
|
||
"question_dense_vector": qa.get("question_dense_vector", []),
|
||
"created_at": created_at
|
||
})
|
||
|
||
# 插入数据
|
||
result = milvus_client.client.insert(
|
||
collection_name=self.collection_name,
|
||
data=data
|
||
)
|
||
|
||
logger.info(f"QA数据入库成功，插入 {len(data)} 条记录")
|
||
return result.get("ids", [])
|
||
|
||
except Exception as e:
|
||
logger.error(f"QA数据入库失败: {str(e)}", exc_info=True)
|
||
return []
|
||
|
||
async def search_qa(self, query_text: List[float], top_k: int = 5) -> List[Dict[str, Any]]:
|
||
"""
|
||
根据查询文本搜索QA数据
|
||
|
||
Args:
|
||
query_text: 查询文本
|
||
top_k: 返回结果数量
|
||
|
||
Returns:
|
||
搜索结果列表，格式 [{"question": "...", "answer": "...", "score": 0.0}, ...]
|
||
"""
|
||
res = milvus_client.client.search(
|
||
collection_name=self.collection_name,
|
||
anns_field="question_dense_vector",
|
||
data=[query_text],
|
||
limit=top_k,
|
||
search_params={"metric_type": "COSINE"},
|
||
output_fields=["question", "answer", "created_at"]
|
||
)
|
||
|
||
score_threshold = 0.6
|
||
filter_results = []
|
||
|
||
# res 是二维结构：[ [单条结果列表] ]
|
||
for hits in res:
|
||
for hit in hits:
|
||
if hit.get("distance",0) >= score_threshold:
|
||
# 按需组装数据
|
||
item = {
|
||
"question": hit.entity.get("question"),
|
||
"answer": hit.entity.get("answer"),
|
||
"created_at": hit.entity.get("created_at"),
|
||
"score": hit.get("distance",0)
|
||
}
|
||
filter_results.append(item)
|
||
|
||
# filter_results 就是最终过滤后的结果
|
||
print("filter_results:", filter_results)
|
||
return filter_results
|
||
# 判空再取值
|
||
# if filter_results:
|
||
# return filter_results[0]
|
||
# else:
|
||
# # 无匹配结果，返回空字典 / None，根据业务选其一
|
||
# return {} |