"""知识库检索:把客户问题向量化后在三个知识集合里检索(只读)。 为什么不复用记忆那套 `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