189 lines
7.4 KiB
Python
189 lines
7.4 KiB
Python
|
"""知识库检索：把客户问题向量化后在三个知识集合里检索（只读）。
|
|||
|
|
|
|||
|
|
为什么不复用记忆那套 `VectorMemoryAdapter`：它只把命中折叠成 `(memory_uuid, score)`，
|
|||
|
|
会把知识块的标题与正文丢掉。而客服回答必须能把**原文与来源**一起交给客户——金融场景
|
|||
|
|
里「答案出自哪份文件的哪一条」本身就是交付物的一部分，丢了正文等于没法给来源引用。
|
|||
|
|
|
|||
|
|
失败语义与基座一致：**任何一步失败都不抛异常给主链路**，而是返回 `degraded=True`
|
|||
|
|
的空结果，由调用方（客服 Agent）据此走「引导客户致电人工客服」的兜底路径。
|
|||
|
|
金融场景下"答不了"是可接受的结果，"答错"不是。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from collections.abc import Awaitable, Callable, Sequence
|
|||
|
|
from dataclasses import dataclass, field
|
|||
|
|
from typing import Any, Protocol
|
|||
|
|
|
|||
|
|
# 三个知识集合（方案 §2.4.4 / §4.1）
|
|||
|
|
FAQ_COLLECTION = "fin_faq_collection"
|
|||
|
|
PRODUCT_COLLECTION = "fin_product_collection"
|
|||
|
|
POLICY_COLLECTION = "fin_policy_collection"
|
|||
|
|
DEFAULT_COLLECTIONS: tuple[str, ...] = (FAQ_COLLECTION, PRODUCT_COLLECTION, POLICY_COLLECTION)
|
|||
|
|
|
|||
|
|
# 检索输出字段：与灌库脚本写入的 schema 对齐
|
|||
|
|
OUTPUT_FIELDS = (
|
|||
|
|
"doc_id", "title", "content", "chapter", "section",
|
|||
|
|
"tags", "doc_no", "version", "source_file", "visibility",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class VectorSearcher(Protocol):
|
|||
|
|
"""只依赖用到的两个方法，便于测试替身注入。"""
|
|||
|
|
|
|||
|
|
def search(self, **kwargs: Any) -> Any: ...
|
|||
|
|
|
|||
|
|
|
|||
|
|
Embedder = Callable[[str], Awaitable[list[float]]]
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass(frozen=True)
|
|||
|
|
class KnowledgeHit:
|
|||
|
|
"""一条知识命中；`score` 为 COSINE 相似度（越大越相似）。"""
|
|||
|
|
|
|||
|
|
doc_id: str
|
|||
|
|
title: str
|
|||
|
|
content: str
|
|||
|
|
score: float
|
|||
|
|
source_file: str = ""
|
|||
|
|
visibility: str = "public"
|
|||
|
|
doc_no: str = ""
|
|||
|
|
version: str = ""
|
|||
|
|
chapter: str = ""
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def reference_title(self) -> str:
|
|||
|
|
"""给客户看的来源标题：优先带内部文件编号，便于人工核对。"""
|
|||
|
|
return f"{self.title}（{self.doc_no}）" if self.doc_no else self.title
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass(frozen=True)
|
|||
|
|
class KnowledgeSearchOutcome:
|
|||
|
|
"""检索结果；`degraded=True` 表示检索链路故障，调用方必须走兜底而非当'没找到'。"""
|
|||
|
|
|
|||
|
|
hits: tuple[KnowledgeHit, ...] = ()
|
|||
|
|
degraded: bool = False
|
|||
|
|
reason: str = ""
|
|||
|
|
searched_collections: tuple[str, ...] = field(default_factory=tuple)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def best(self) -> KnowledgeHit | None:
|
|||
|
|
return self.hits[0] if self.hits else None
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def top_score(self) -> float:
|
|||
|
|
return self.hits[0].score if self.hits else 0.0
|
|||
|
|
|
|||
|
|
|
|||
|
|
class KnowledgeSearchService:
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
client: VectorSearcher | None,
|
|||
|
|
embedder: Embedder | None,
|
|||
|
|
*,
|
|||
|
|
collections: Sequence[str] = DEFAULT_COLLECTIONS,
|
|||
|
|
) -> None:
|
|||
|
|
self._client = client
|
|||
|
|
self._embedder = embedder
|
|||
|
|
self._collections = tuple(collections)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def available(self) -> bool:
|
|||
|
|
"""向量库与向量化能力是否都在位；缺任一项都不做检索，直接走兜底。"""
|
|||
|
|
return self._client is not None and self._embedder is not None
|
|||
|
|
|
|||
|
|
async def search(
|
|||
|
|
self,
|
|||
|
|
query: str,
|
|||
|
|
*,
|
|||
|
|
collections: Sequence[str] | None = None,
|
|||
|
|
top_k: int = 5,
|
|||
|
|
include_internal: bool = False,
|
|||
|
|
) -> KnowledgeSearchOutcome:
|
|||
|
|
"""检索知识库。
|
|||
|
|
|
|||
|
|
`include_internal=False`（默认）时在 Milvus 侧就过滤掉 `visibility=internal` 的块，
|
|||
|
|
内部资料不进入面向客户的答案——这是检索层的硬隔离，不依赖提示词约束。
|
|||
|
|
"""
|
|||
|
|
text = query.strip()
|
|||
|
|
if not text:
|
|||
|
|
return KnowledgeSearchOutcome(reason="empty_query")
|
|||
|
|
if not self.available:
|
|||
|
|
return KnowledgeSearchOutcome(degraded=True, reason="vector_backend_unavailable")
|
|||
|
|
|
|||
|
|
client, embedder = self._client, self._embedder
|
|||
|
|
if client is None or embedder is None:
|
|||
|
|
# 与 available 重复，但这里需要类型收窄（mypy 不跨属性判断 Optional）
|
|||
|
|
return KnowledgeSearchOutcome(degraded=True, reason="vector_backend_unavailable")
|
|||
|
|
try:
|
|||
|
|
vector = await embedder(text)
|
|||
|
|
except Exception:
|
|||
|
|
return KnowledgeSearchOutcome(degraded=True, reason="embedding_failed")
|
|||
|
|
if not vector:
|
|||
|
|
return KnowledgeSearchOutcome(degraded=True, reason="embedding_empty")
|
|||
|
|
|
|||
|
|
targets = tuple(collections or self._collections)
|
|||
|
|
expression = None if include_internal else 'visibility == "public"'
|
|||
|
|
collected: list[KnowledgeHit] = []
|
|||
|
|
failures = 0
|
|||
|
|
for collection in targets:
|
|||
|
|
try:
|
|||
|
|
raw = client.search(
|
|||
|
|
collection_name=collection,
|
|||
|
|
data=[vector],
|
|||
|
|
limit=max(1, min(top_k, 20)),
|
|||
|
|
output_fields=list(OUTPUT_FIELDS),
|
|||
|
|
filter=expression,
|
|||
|
|
)
|
|||
|
|
except Exception:
|
|||
|
|
failures += 1
|
|||
|
|
continue
|
|||
|
|
collected.extend(self._parse(raw, collection))
|
|||
|
|
|
|||
|
|
if not collected and failures == len(targets) and targets:
|
|||
|
|
# 三个集合全查失败：是链路故障，不是"知识库里没有"
|
|||
|
|
return KnowledgeSearchOutcome(
|
|||
|
|
degraded=True, reason="search_failed", searched_collections=targets
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
collected.sort(key=lambda hit: hit.score, reverse=True)
|
|||
|
|
# 同一内容可能同时存在于产品手册与问答对里，按 doc_id 去重保留最高分
|
|||
|
|
deduped: list[KnowledgeHit] = []
|
|||
|
|
seen: set[str] = set()
|
|||
|
|
for hit in collected:
|
|||
|
|
if hit.doc_id in seen:
|
|||
|
|
continue
|
|||
|
|
seen.add(hit.doc_id)
|
|||
|
|
deduped.append(hit)
|
|||
|
|
return KnowledgeSearchOutcome(
|
|||
|
|
hits=tuple(deduped[: max(1, top_k)]),
|
|||
|
|
degraded=failures > 0,
|
|||
|
|
reason="partial_collection_failure" if failures else "",
|
|||
|
|
searched_collections=targets,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _parse(raw: Any, collection: str) -> list[KnowledgeHit]:
|
|||
|
|
"""把 pymilvus 的 `[[{id, distance, entity}]]` 折叠成命中列表（纯函数，不抛异常）。"""
|
|||
|
|
hits: list[KnowledgeHit] = []
|
|||
|
|
groups = raw if isinstance(raw, (list, tuple)) else [raw]
|
|||
|
|
for group in groups:
|
|||
|
|
rows = group if isinstance(group, (list, tuple)) else [group]
|
|||
|
|
for row in rows:
|
|||
|
|
entity = row.get("entity") if isinstance(row, dict) else None
|
|||
|
|
if not isinstance(entity, dict):
|
|||
|
|
continue
|
|||
|
|
content = str(entity.get("content") or "")
|
|||
|
|
if not content:
|
|||
|
|
continue # 没有正文的命中无法作为答案来源，直接丢弃而不是猜造
|
|||
|
|
hits.append(KnowledgeHit(
|
|||
|
|
doc_id=str(entity.get("doc_id") or ""),
|
|||
|
|
title=str(entity.get("title") or ""),
|
|||
|
|
content=content,
|
|||
|
|
score=float(row.get("distance") or 0.0),
|
|||
|
|
source_file=str(entity.get("source_file") or ""),
|
|||
|
|
visibility=str(entity.get("visibility") or "public"),
|
|||
|
|
doc_no=str(entity.get("doc_no") or ""),
|
|||
|
|
version=str(entity.get("version") or ""),
|
|||
|
|
chapter=str(entity.get("chapter") or ""),
|
|||
|
|
))
|
|||
|
|
return hits
|