67 lines
3.2 KiB
Python
67 lines
3.2 KiB
Python
from typing import Any, Protocol
|
|
|
|
from app.core.contracts import RequestContext
|
|
from app.core.errors import RecoverableAgentError
|
|
from app.core.knowledge_contracts import KnowledgeHit, KnowledgeQuery, KnowledgeSearchResult
|
|
from app.service.knowledge_config import KnowledgeRuntimeConfig
|
|
|
|
# ruff: noqa: E501
|
|
|
|
|
|
class KnowledgeEmbedder(Protocol):
|
|
async def embed(self, text: str) -> list[float]: ...
|
|
|
|
|
|
class KnowledgeVectorStore(Protocol):
|
|
async def search(self, collection: str, vector: list[float], top_k: int) -> list[dict[str, Any]]: ...
|
|
|
|
|
|
class KnowledgeAuthority(Protocol):
|
|
async def filter_published(self, hits: tuple[KnowledgeHit, ...]) -> list[KnowledgeHit]: ...
|
|
|
|
async def search_keyword(
|
|
self, query: KnowledgeQuery, collections: tuple[str, ...], top_k: int
|
|
) -> list[KnowledgeHit]: ...
|
|
|
|
|
|
class KnowledgeRetrievalService:
|
|
def __init__(
|
|
self, embedder: KnowledgeEmbedder, vector_store: KnowledgeVectorStore,
|
|
config: KnowledgeRuntimeConfig, authority: KnowledgeAuthority,
|
|
) -> None:
|
|
self._embedder = embedder
|
|
self._vector_store = vector_store
|
|
self._config = config
|
|
self._authority = authority
|
|
|
|
async def search(self, query: KnowledgeQuery, context: RequestContext) -> KnowledgeSearchResult:
|
|
del context
|
|
targets = tuple(self._config.routes[intent] for intent in query.intents if intent in self._config.routes)
|
|
if not targets:
|
|
return KnowledgeSearchResult()
|
|
collections: list[str] = []
|
|
candidates: list[KnowledgeHit] = []
|
|
try:
|
|
vector = await self._embedder.embed(query.query)
|
|
if len(vector) != self._config.vector_dim:
|
|
raise RecoverableAgentError("嵌入维度与集合定义不一致")
|
|
for collection, configured_top_k in targets:
|
|
collections.append(collection)
|
|
for raw in await self._vector_store.search(
|
|
collection, vector, min(query.top_k, configured_top_k)
|
|
):
|
|
knowledge_id, snippet, score = raw.get("knowledge_id"), raw.get("snippet"), raw.get("score")
|
|
if isinstance(knowledge_id, str) and isinstance(snippet, str) and isinstance(score, (int, float)):
|
|
if not isinstance(score, bool) and self._config.similarity_threshold <= score <= 1:
|
|
candidates.append(KnowledgeHit(knowledge_id=knowledge_id, collection=collection, snippet=snippet, score=float(score)))
|
|
except RecoverableAgentError:
|
|
fallback_collections = tuple(dict.fromkeys(collection for collection, _ in targets))
|
|
fallback_top_k = max(min(query.top_k, configured_top_k) for _, configured_top_k in targets)
|
|
fallback_hits = await self._authority.search_keyword(query, fallback_collections, fallback_top_k)
|
|
return KnowledgeSearchResult(
|
|
hits=tuple(fallback_hits), degraded=True, degradation_reason="milvus_unavailable",
|
|
searched_collections=fallback_collections,
|
|
)
|
|
hits = await self._authority.filter_published(tuple(candidates))
|
|
return KnowledgeSearchResult(hits=tuple(hits), searched_collections=tuple(dict.fromkeys(collections)))
|