101 lines
3.3 KiB
Python
101 lines
3.3 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 [] |