355 lines
14 KiB
Python
355 lines
14 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 IntentClassificationRepository:
|
||
collection_name: str = 'intent_classification'
|
||
|
||
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:
|
||
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="意图分类向量集合"
|
||
)
|
||
|
||
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(field_name="intent_category", datatype=DataType.VARCHAR, max_length=128)
|
||
schema.add_field(field_name="content", datatype=DataType.VARCHAR, max_length=65535)
|
||
schema.add_field(field_name="review_content", datatype=DataType.VARCHAR, max_length=65535)
|
||
schema.add_field(field_name="content_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="content_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
|
||
# 迁移
|
||
def migrate_add_review_content(self) -> bool:
|
||
if not milvus_client.has_collection(self.collection_name):
|
||
logger.error(f"集合 {self.collection_name} 不存在，无法迁移")
|
||
return False
|
||
|
||
temp_collection = f"{self.collection_name}_temp"
|
||
|
||
if milvus_client.has_collection(temp_collection):
|
||
milvus_client.client.drop_collection(temp_collection)
|
||
|
||
schema = milvus_client.client.create_schema(
|
||
auto_id=True,
|
||
enable_dynamic_field=True,
|
||
description="意图分类向量集合（迁移：新增review_content）"
|
||
)
|
||
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
|
||
schema.add_field(field_name="intent_category", datatype=DataType.VARCHAR, max_length=128)
|
||
schema.add_field(field_name="content", datatype=DataType.VARCHAR, max_length=65535)
|
||
schema.add_field(field_name="review_content", datatype=DataType.VARCHAR, max_length=65535)
|
||
schema.add_field(field_name="content_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="content_vector",
|
||
index_type="HNSW",
|
||
metric_type="COSINE",
|
||
params={"M": 32, "efConstruction": 300}
|
||
)
|
||
milvus_client.client.create_collection(
|
||
collection_name=temp_collection,
|
||
schema=schema,
|
||
index_params=index_params
|
||
)
|
||
logger.info(f"创建临时集合 {temp_collection} 成功")
|
||
|
||
all_data = []
|
||
offset = 0
|
||
batch_size = 1000
|
||
while True:
|
||
batch = milvus_client.client.query(
|
||
collection_name=self.collection_name,
|
||
output_fields=["intent_category", "content", "content_vector", "created_at"],
|
||
limit=batch_size,
|
||
offset=offset
|
||
)
|
||
if not batch:
|
||
break
|
||
all_data.extend(batch)
|
||
if len(batch) < batch_size:
|
||
break
|
||
offset += batch_size
|
||
|
||
logger.info(f"查询到旧数据 {len(all_data)} 条，开始迁移")
|
||
|
||
migrated = 0
|
||
for i in range(0, len(all_data), batch_size):
|
||
batch = all_data[i:i + batch_size]
|
||
insert_batch = []
|
||
for item in batch:
|
||
insert_batch.append({
|
||
"intent_category": item.get("intent_category", ""),
|
||
"content": item.get("content", ""),
|
||
"review_content": "",
|
||
"content_vector": item.get("content_vector", []),
|
||
"created_at": item.get("created_at", 0)
|
||
})
|
||
milvus_client.client.insert(
|
||
collection_name=temp_collection,
|
||
data=insert_batch
|
||
)
|
||
migrated += len(insert_batch)
|
||
logger.info(f"迁移进度: {migrated}/{len(all_data)}")
|
||
|
||
logger.info(f"数据迁移完成，共 {len(all_data)} 条")
|
||
|
||
milvus_client.client.drop_collection(self.collection_name)
|
||
logger.info(f"旧集合 {self.collection_name} 已删除")
|
||
|
||
milvus_client.client.rename_collection(temp_collection, self.collection_name)
|
||
logger.info(f"临时集合重命名为 {self.collection_name}，迁移完成")
|
||
|
||
return True
|
||
|
||
async def save_intent_data(self, data: List[Dict[str, Any]]) -> List[int]:
|
||
try:
|
||
created_at = int(datetime.now().timestamp())
|
||
insert_data = []
|
||
for item in data:
|
||
insert_data.append({
|
||
"intent_category": item.get("intent_category", ""),
|
||
"content": item.get("content", ""),
|
||
"review_content": item.get("review_content", ""),
|
||
"content_vector": item.get("content_vector", []),
|
||
"created_at": created_at
|
||
})
|
||
|
||
result = milvus_client.client.insert(
|
||
collection_name=self.collection_name,
|
||
data=insert_data
|
||
)
|
||
|
||
logger.info(f"意图分类数据入库成功，插入 {len(data)} 条记录")
|
||
return result.get("ids", [])
|
||
|
||
except Exception as e:
|
||
logger.error(f"意图分类数据入库失败: {str(e)}", exc_info=True)
|
||
return []
|
||
|
||
async def upsert_intent_data(self, data: List[Dict[str, Any]]) -> List[int]:
|
||
try:
|
||
created_at = int(datetime.now().timestamp())
|
||
all_ids = []
|
||
for item in data:
|
||
category = item.get("intent_category", "")
|
||
content = item.get("content", "")
|
||
vector = item.get("content_vector", [])
|
||
|
||
escaped_content = content.replace("\\", "\\\\").replace('"', '\\"')
|
||
expr = f'intent_category == "{category}"'
|
||
existing = milvus_client.client.query(
|
||
collection_name=self.collection_name,
|
||
filter=expr,
|
||
output_fields=["id"],
|
||
limit=1000
|
||
)
|
||
if existing:
|
||
existing_ids = [e["id"] for e in existing]
|
||
milvus_client.client.delete(
|
||
collection_name=self.collection_name,
|
||
ids=existing_ids
|
||
)
|
||
logger.info(f"删除旧记录 {len(existing_ids)} 条: category={category}")
|
||
|
||
insert_result = milvus_client.client.insert(
|
||
collection_name=self.collection_name,
|
||
data=[{
|
||
"intent_category": category,
|
||
"content": content,
|
||
"review_content": item.get("review_content", ""),
|
||
"content_vector": vector,
|
||
"created_at": created_at
|
||
}]
|
||
)
|
||
all_ids.extend(insert_result.get("ids", []))
|
||
|
||
logger.info(f"意图分类数据upsert完成，共处理 {len(data)} 条")
|
||
return all_ids
|
||
|
||
except Exception as e:
|
||
logger.error(f"意图分类数据upsert失败: {str(e)}", exc_info=True)
|
||
return []
|
||
|
||
async def search_by_intent_category(self, intent_category: List[str], top_k: int = 100) -> List[Dict[str, Any]]:
|
||
if isinstance(intent_category, str):
|
||
intent_category = [intent_category]
|
||
categories_str = ", ".join(f'"{c}"' for c in intent_category)
|
||
expr = f'intent_category in [{categories_str}]'
|
||
result = milvus_client.client.query(
|
||
collection_name=self.collection_name,
|
||
filter=expr,
|
||
output_fields=["intent_category", "content", "review_content", "created_at"],
|
||
limit=top_k
|
||
)
|
||
return result
|
||
|
||
async def search_by_vector(self, query_vector: List[float], top_k: int = 5) -> List[Dict[str, Any]]:
|
||
res = milvus_client.client.search(
|
||
collection_name=self.collection_name,
|
||
anns_field="content_vector",
|
||
data=[query_vector],
|
||
limit=top_k,
|
||
search_params={
|
||
"metric_type": "COSINE",
|
||
"efSearch": 300
|
||
},
|
||
output_fields=["intent_category", "content", "review_content", "created_at"]
|
||
)
|
||
|
||
results = []
|
||
for hits in res:
|
||
for hit in hits:
|
||
item = {
|
||
"intent_category": hit.entity.get("intent_category"),
|
||
"content": hit.entity.get("content"),
|
||
"review_content": hit.entity.get("review_content"),
|
||
"created_at": hit.entity.get("created_at"),
|
||
"score": hit.get("distance", 0)
|
||
}
|
||
results.append(item)
|
||
return results
|
||
|
||
async def get_all_intent_categories(self) -> List[str]:
|
||
result = milvus_client.client.query(
|
||
collection_name=self.collection_name,
|
||
output_fields=["intent_category"],
|
||
limit=1000
|
||
)
|
||
categories = list(set([item.get("intent_category", "") for item in result]))
|
||
return categories
|
||
|
||
async def get_style_intent_categories(self) -> List[str]:
|
||
all_categories = await self.get_all_intent_categories()
|
||
return [c for c in all_categories]
|
||
|
||
async def get_content_by_intent_category(self, intent_category: str) -> List[Dict[str, str]]:
|
||
result = milvus_client.client.query(
|
||
collection_name=self.collection_name,
|
||
filter=f'intent_category == "{intent_category}"',
|
||
output_fields=["content", "review_content"],
|
||
limit=1000
|
||
)
|
||
return [
|
||
{
|
||
"content": item.get("content", ""),
|
||
"review_content": item.get("review_content", "")
|
||
}
|
||
for item in result if item.get("content")
|
||
]
|
||
|
||
async def delete_today_data(self) -> int:
|
||
today_start = int(datetime.now().replace(hour=0, minute=0, second=0, microsecond=0).timestamp())
|
||
query_result = milvus_client.client.query(
|
||
collection_name=self.collection_name,
|
||
filter=f'created_at >= {today_start}',
|
||
output_fields=["id"],
|
||
limit=10000
|
||
)
|
||
if not query_result:
|
||
logger.info("今日无数据可删除")
|
||
return 0
|
||
ids_to_delete = [item["id"] for item in query_result]
|
||
milvus_client.client.delete(
|
||
collection_name=self.collection_name,
|
||
ids=ids_to_delete
|
||
)
|
||
logger.info(f"删除今日数据 {len(ids_to_delete)} 条")
|
||
return len(ids_to_delete)
|
||
|
||
def sync_content_to_review_content(self) -> int:
|
||
result = milvus_client.client.query(
|
||
collection_name=self.collection_name,
|
||
output_fields=["id", "intent_category", "content", "review_content", "content_vector", "created_at"],
|
||
limit=10000
|
||
)
|
||
if not result:
|
||
logger.info("无数据需要同步")
|
||
return 0
|
||
|
||
updated = 0
|
||
batch_size = 1000
|
||
for i in range(0, len(result), batch_size):
|
||
batch = result[i:i + batch_size]
|
||
ids_to_delete = []
|
||
insert_batch = []
|
||
for item in batch:
|
||
ids_to_delete.append(item["id"])
|
||
insert_batch.append({
|
||
"intent_category": item.get("intent_category", ""),
|
||
"content": item.get("content", ""),
|
||
"review_content": item.get("content", ""),
|
||
"content_vector": item.get("content_vector", []),
|
||
"created_at": item.get("created_at", 0)
|
||
})
|
||
milvus_client.client.delete(
|
||
collection_name=self.collection_name,
|
||
ids=ids_to_delete
|
||
)
|
||
milvus_client.client.insert(
|
||
collection_name=self.collection_name,
|
||
data=insert_batch
|
||
)
|
||
updated += len(insert_batch)
|
||
logger.info(f"同步进度: {updated}/{len(result)}")
|
||
|
||
logger.info(f"content → review_content 同步完成，共更新 {updated} 条")
|
||
return updated
|
||
|
||
|
||
async def main():
|
||
milvus_client.init()
|
||
repo = IntentClassificationRepository()
|
||
|
||
# 1. 首次部署：迁移新增 review_content 字段
|
||
# success = repo.migrate_add_review_content()
|
||
# print(f"迁移结果: {success}")
|
||
|
||
# 2. 将 content 同步写入 review_content
|
||
# count = repo.sync_content_to_review_content()
|
||
# print(f"同步完成，共 {count} 条")
|
||
|
||
style_categories = await repo.get_style_intent_categories()
|
||
print(style_categories)
|
||
|
||
if __name__ == "__main__":
|
||
import asyncio
|
||
asyncio.run(main())
|
||
|