129 lines
4.7 KiB
Python
129 lines
4.7 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 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 if "风格" in c] |