""" HybridRetriever: 混合检索器 策略：pgvector 向量搜索 + PostgreSQL ILIKE 关键词搜索 + RRF（Reciprocal Rank Fusion）融合排序 """ import logging import uuid from typing import TYPE_CHECKING from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession if TYPE_CHECKING: from .embedder import EmbeddingService logger = logging.getLogger(__name__) class HybridRetriever: """混合检索器：向量搜索 + 关键词搜索 + RRF 融合""" def __init__( self, embedder: "EmbeddingService", vector_weight: float = 0.7, keyword_weight: float = 0.3, ): self.embedder = embedder self.vector_weight = vector_weight self.keyword_weight = keyword_weight # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ async def search( self, session: AsyncSession, query: str, knowledge_base_ids: list[str], top_k: int = 10, ) -> list[dict]: """ 混合检索： 1. 获取 query embedding 2. 向量搜索（pgvector cosine distance） 3. 关键词搜索（ILIKE） 4. RRF 融合排序 Returns: list of dicts: { "chunk_id": str, "content": str, "score": float, "document_id": str, "document_title": str, "metadata": dict, } """ if not query or not knowledge_base_ids: return [] # 并发执行向量检索和关键词检索 # （asyncio 环境中顺序 await 即可，DB session 不支持真并发） query_embedding = await self.embedder.embed(query) vector_results = await self._vector_search( session, query_embedding, knowledge_base_ids, top_k * 2 ) keyword_results = await self._keyword_search( session, query, knowledge_base_ids, top_k * 2 ) fused = self._rrf_fusion(vector_results, keyword_results) return fused[:top_k] # ------------------------------------------------------------------ # Internal: vector search # ------------------------------------------------------------------ async def _vector_search( self, session: AsyncSession, query_embedding: list[float], kb_ids: list[str], top_k: int, ) -> list[dict]: """ pgvector 向量搜索（cosine distance <=>）。 使用原生 SQL 以支持 pgvector 操作符。 """ # 将 embedding list 转为 pgvector 字符串格式 embedding_str = "[" + ",".join(str(v) for v in query_embedding) + "]" # kb_ids 转为 UUID 数组，兼容 PostgreSQL ANY() kb_uuid_strs = [str(uid) for uid in kb_ids] sql = text( """ SELECT kc.id AS chunk_id, kc.content AS content, kd.id AS document_id, kd.title AS document_title, kc.metadata AS metadata, 1 - (kc.embedding <=> CAST(:query_vec AS vector)) AS score FROM knowledge_chunks kc JOIN knowledge_documents kd ON kc.document_id = kd.id WHERE kd.knowledge_base_id = ANY(CAST(:kb_ids AS uuid[])) AND kd.status = 'ready' AND kc.embedding IS NOT NULL ORDER BY kc.embedding <=> CAST(:query_vec AS vector) LIMIT :top_k """ ) try: result = await session.execute( sql, { "query_vec": embedding_str, "kb_ids": "{" + ",".join(kb_uuid_strs) + "}", "top_k": top_k, }, ) rows = result.mappings().all() except Exception as e: logger.warning(f"Vector search failed: {e}; falling back to empty results") return [] return [ { "chunk_id": str(row["chunk_id"]), "content": row["content"], "score": float(row["score"]) if row["score"] is not None else 0.0, "document_id": str(row["document_id"]), "document_title": row["document_title"], "metadata": row["metadata"] or {}, } for row in rows ] # ------------------------------------------------------------------ # Internal: keyword search # ------------------------------------------------------------------ async def _keyword_search( self, session: AsyncSession, query: str, kb_ids: list[str], top_k: int, ) -> list[dict]: """ 关键词搜索（ILIKE 模糊匹配）。 后续可升级为 ts_vector 全文检索。 """ kb_uuid_strs = [str(uid) for uid in kb_ids] like_pattern = f"%{query}%" sql = text( """ SELECT kc.id AS chunk_id, kc.content AS content, kd.id AS document_id, kd.title AS document_title, kc.metadata AS metadata, 1.0 AS score FROM knowledge_chunks kc JOIN knowledge_documents kd ON kc.document_id = kd.id WHERE kd.knowledge_base_id = ANY(CAST(:kb_ids AS uuid[])) AND kd.status = 'ready' AND kc.content ILIKE :pattern LIMIT :top_k """ ) try: result = await session.execute( sql, { "kb_ids": "{" + ",".join(kb_uuid_strs) + "}", "pattern": like_pattern, "top_k": top_k, }, ) rows = result.mappings().all() except Exception as e: logger.warning(f"Keyword search failed: {e}; falling back to empty results") return [] return [ { "chunk_id": str(row["chunk_id"]), "content": row["content"], "score": 1.0, "document_id": str(row["document_id"]), "document_title": row["document_title"], "metadata": row["metadata"] or {}, } for row in rows ] # ------------------------------------------------------------------ # Internal: RRF fusion # ------------------------------------------------------------------ def _rrf_fusion( self, vector_results: list[dict], keyword_results: list[dict], k: int = 60, ) -> list[dict]: """ Reciprocal Rank Fusion (RRF)。 score = Σ ( weight_i / (k + rank_i) ) 对所有结果列表求和 """ scores: dict[str, float] = {} chunk_data: dict[str, dict] = {} def _accumulate(results: list[dict], weight: float) -> None: for rank, item in enumerate(results, start=1): cid = item["chunk_id"] rrf_score = weight / (k + rank) scores[cid] = scores.get(cid, 0.0) + rrf_score if cid not in chunk_data: chunk_data[cid] = item _accumulate(vector_results, self.vector_weight) _accumulate(keyword_results, self.keyword_weight) # 按 RRF 分数降序排列 sorted_ids = sorted(scores, key=lambda cid: scores[cid], reverse=True) fused = [] for cid in sorted_ids: item = dict(chunk_data[cid]) item["score"] = round(scores[cid], 6) fused.append(item) return fused