zhanghongyu_0626
2945108f66
- Added new configuration for knowledge base root directory in `.env.example` and `settings.py`. - Implemented `find_products` method in `CoreReadOnlyRepository` for fuzzy product search based on user queries. - Introduced `search_cs_knowledge` function in `rag_service.py` to facilitate semantic search across new `fin_*` collections. - Updated document parsing to support Markdown and YAML front-matter for knowledge base entries. - Created multiple new FAQ and policy documents in the `data/kb_collections` directory to enrich the knowledge base. This update significantly improves the knowledge retrieval capabilities for customer service interactions, ensuring more relevant and accurate responses.
104 lines
3.5 KiB
Python
104 lines
3.5 KiB
Python
"""Milvus 封装：Milvus Lite 连接、fin_* Collection CRUD、语义检索（客服 KB 第二套）。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from app.config.settings import settings
|
|
from app.service import milvus_service
|
|
from pymilvus import DataType # noqa: E402 milvus_service 已做 MILVUS_URI 环境防御
|
|
|
|
# build_collections.py 与检索共用
|
|
COLLECTION_FIELDS: dict[str, list[str]] = {
|
|
"fin_faq": ["question", "answer", "category", "source_doc", "chunk_no", "chunk_text"],
|
|
"fin_product": ["product_name", "risk_level", "doc_type", "source_doc", "chunk_no", "chunk_text"],
|
|
"fin_policy": ["policy_name", "chapter", "source_doc", "chunk_no", "chunk_text"],
|
|
}
|
|
|
|
_VARCHAR_MAX = {
|
|
"question": 512,
|
|
"answer": 2048,
|
|
"category": 64,
|
|
"product_name": 256,
|
|
"risk_level": 8,
|
|
"doc_type": 32,
|
|
"policy_name": 256,
|
|
"chapter": 128,
|
|
"source_doc": 128,
|
|
"chunk_text": 8192,
|
|
}
|
|
|
|
|
|
def _schema_for(collection: str, dim: int):
|
|
fields = COLLECTION_FIELDS[collection]
|
|
schema = milvus_service.MilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
|
|
schema.add_field("id", DataType.VARCHAR, is_primary=True, max_length=64)
|
|
schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=dim)
|
|
for name in fields:
|
|
if name == "chunk_no":
|
|
schema.add_field(name, DataType.INT64)
|
|
else:
|
|
schema.add_field(name, DataType.VARCHAR, max_length=_VARCHAR_MAX.get(name, 256))
|
|
return schema
|
|
|
|
|
|
class MilvusKbClient:
|
|
"""fin_faq / fin_product / fin_policy 三库客户端。"""
|
|
|
|
def __init__(self, client: milvus_service.MilvusClient | None = None) -> None:
|
|
self._client = client or milvus_service.milvus_client()
|
|
|
|
def ensure_collection(self, name: str) -> None:
|
|
if name not in COLLECTION_FIELDS:
|
|
raise ValueError(f"unknown collection: {name}")
|
|
dim = settings.embed_dim
|
|
if self._client.has_collection(name):
|
|
self._client.load_collection(name)
|
|
return
|
|
index = self._client.prepare_index_params()
|
|
index.add_index(field_name="embedding", index_type="AUTOINDEX", metric_type="COSINE")
|
|
self._client.create_collection(name, schema=_schema_for(name, dim), index_params=index)
|
|
self._client.load_collection(name)
|
|
|
|
def drop_collection(self, name: str) -> None:
|
|
if self._client.has_collection(name):
|
|
self._client.drop_collection(name)
|
|
|
|
def insert(self, name: str, records: list[dict[str, Any]]) -> int:
|
|
if not records:
|
|
return 0
|
|
self.ensure_collection(name)
|
|
self._client.upsert(name, data=records)
|
|
return len(records)
|
|
|
|
def search(self, name: str, vector: list[float], top_k: int = 3) -> list[dict[str, Any]]:
|
|
if not self._client.has_collection(name):
|
|
return []
|
|
self._client.load_collection(name)
|
|
output_fields = list(COLLECTION_FIELDS[name])
|
|
results = self._client.search(
|
|
name,
|
|
data=[vector],
|
|
limit=top_k,
|
|
output_fields=output_fields,
|
|
)
|
|
hits = results[0] if results else []
|
|
out: list[dict[str, Any]] = []
|
|
for hit in hits:
|
|
entity = hit.get("entity", {})
|
|
out.append(
|
|
{
|
|
"id": hit.get("id"),
|
|
"score": float(hit.get("distance", 0.0)),
|
|
**entity,
|
|
}
|
|
)
|
|
return out
|
|
|
|
def close(self) -> None:
|
|
self._client.close()
|
|
|
|
|
|
def get_milvus_client() -> MilvusKbClient:
|
|
return MilvusKbClient()
|