215 lines
8.0 KiB
Python
215 lines
8.0 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="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
|
||
|
||
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", ""),
|
||
"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,
|
||
"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", "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", "created_at"]
|
||
)
|
||
|
||
results = []
|
||
for hits in res:
|
||
for hit in hits:
|
||
item = {
|
||
"intent_category": hit.entity.get("intent_category"),
|
||
"content": hit.entity.get("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[str]:
|
||
result = milvus_client.client.query(
|
||
collection_name=self.collection_name,
|
||
filter=f'intent_category == "{intent_category}"',
|
||
output_fields=["content"],
|
||
limit=1000
|
||
)
|
||
return [item.get("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)
|
||
|
||
|
||
async def main():
|
||
milvus_client.init()
|
||
repo = IntentClassificationRepository()
|
||
style_categories = await repo.delete_today_data()
|
||
print(style_categories)
|
||
|
||
if __name__ == "__main__":
|
||
# today_start = int(datetime.now().replace(hour=0, minute=0, second=0, microsecond=0).timestamp())
|
||
# print(today_start)
|
||
import asyncio
|
||
asyncio.run(main())
|
||
|