"""双索引检索引擎 — embedding/keywords/blend 三模式。 - embedding: pgvector 语义检索（LlamaIndex PGVectorStore.aquery） - keywords: PG 全文检索（jieba 分词 + tsquery_rank） - blend: 并行执行两种检索，按分数归一化后合并去重排序 参考 LlamaIndex hybrid retriever 模式（VectorStoreRetriever + 全文检索）。 ACL 过滤由调用方在传入 `kb_ids` 前完成（见 acl.filter_kb_by_user_acl）。 """ from __future__ import annotations import asyncio import logging from typing import TYPE_CHECKING from sqlalchemy import text from agentkit.rag_platform.fulltext import KB_CHUNKS_TABLE, build_tsquery from agentkit.rag_platform.models import QueryMode, QueryResult if TYPE_CHECKING: from llama_index.core.embeddings import BaseEmbedding from llama_index.core.vector_stores.types import VectorStoreQuery from llama_index.vector_stores.postgres import PGVectorStore logger = logging.getLogger(__name__) # ponytail: 默认归一化策略 — min-max，将每种检索的分数映射到 [0,1] # 升级路径：若需更精细的融合权重，可引入 RRF (Reciprocal Rank Fusion) 或学习权重 _BLEND_WEIGHT_EMBEDDING = 0.6 _BLEND_WEIGHT_KEYWORDS = 0.4 def _normalize_scores(scores: list[float]) -> list[float]: """min-max 归一化分数到 [0, 1]。 空列表返回空。常数列表（含单元素）返回全 1.0 — 所有结果等价于最高相关度。 """ if not scores: return [] lo, hi = min(scores), max(scores) if hi - lo < 1e-9: return [1.0 for _ in scores] return [(s - lo) / (hi - lo) for s in scores] class RetrievalEngine: """双索引检索引擎 — embedding/keywords/blend 三模式。 Args: vector_store: LlamaIndex PGVectorStore 实例（embedding 模式使用） session_factory: SQLAlchemy async session factory（keywords 模式使用） embed_model: LlamaIndex BaseEmbedding（embedding 模式使用） """ def __init__( self, vector_store: "PGVectorStore", session_factory, embed_model: "BaseEmbedding | None" = None, ) -> None: self._vector_store = vector_store self._sf = session_factory self._embed_model = embed_model async def retrieve( self, query: str, kb_ids: list[str], mode: QueryMode, top_k: int = 5, ) -> list[QueryResult]: """检索入口 — 根据 mode 分发到具体实现。 Args: query: 查询文本 kb_ids: 限定检索的知识库 ID 列表（已通过 ACL 过滤） mode: 检索模式 top_k: 返回结果数 Returns: QueryResult 列表（按分数降序）。无结果时返回空列表。 """ if not kb_ids: return [] if mode == QueryMode.embedding: return await self._retrieve_embedding(query, kb_ids, top_k) elif mode == QueryMode.keywords: return await self._retrieve_keywords(query, kb_ids, top_k) elif mode == QueryMode.blend: return await self._retrieve_blend(query, kb_ids, top_k) else: # pragma: no cover — 枚举已穷尽 raise ValueError(f"Unsupported query mode: {mode}") async def _retrieve_embedding( self, query: str, kb_ids: list[str], top_k: int, ) -> list[QueryResult]: """pgvector 语义检索。 使用 LlamaIndex vector_store.aquery 执行向量检索，结果按 metadata.kb_id 过滤到 kb_ids 子集。 """ if self._embed_model is None: raise ValueError("embed_model is required for embedding retrieval mode") from llama_index.core.vector_stores.types import VectorStoreQuery query_embedding = await self._embed_model.aget_text_embedding(query) vs_query: VectorStoreQuery = VectorStoreQuery( query_embedding=query_embedding, similarity_top_k=top_k, ) result = await self._vector_store.aquery(vs_query) kb_set = set(kb_ids) out: list[QueryResult] = [] for node, score in zip(result.nodes, result.similarities): meta = node.metadata or {} node_kb = meta.get("kb_id", "") if node_kb not in kb_set: continue out.append( QueryResult( chunk_id=node.node_id, content=node.get_content(), score=float(score) if score is not None else 0.0, metadata=meta, document_id=meta.get("document_id", ""), kb_id=node_kb, ) ) if len(out) >= top_k: break return out async def _retrieve_keywords( self, query: str, kb_ids: list[str], top_k: int, ) -> list[QueryResult]: """PG 全文检索（jieba 分词 + tsquery）。 使用 `ts_rank(search_vector, tsquery)` 排序，按 kb_id 过滤。 kb_id 存储在 metadata_ JSON 列中（LlamaIndex PGVectorStore schema）， 用 `metadata_->>'kb_id'` 提取。 """ tsquery = build_tsquery(query) if not tsquery: return [] # ponytail: 用 ANY(%s) 传 kb_ids 列表，避免字符串拼接 # 升级路径：若 kb_ids 数量超过 PG 参数限制（32k），需分批查询 # 列名参考 LlamaIndex PGVectorStore 默认 schema： # id / embedding / text / metadata_ (JSON) / document_id / search_vector sql = text( f""" SELECT id, text, metadata_, document_id, ts_rank(search_vector, to_tsquery('simple', :tsquery)) AS score FROM {KB_CHUNKS_TABLE} WHERE search_vector @@ to_tsquery('simple', :tsquery) AND metadata_->>'kb_id' = ANY(:kb_ids) ORDER BY score DESC LIMIT :top_k """ # noqa: S608 — 表名为常量 ) async with self._sf() as db: result = await db.execute( sql, { "tsquery": tsquery, "kb_ids": list(kb_ids), "top_k": top_k, }, ) rows = result.all() out: list[QueryResult] = [] for row in rows: chunk_id, content, metadata, document_id, score = row meta: dict[str, object] = metadata if isinstance(metadata, dict) else {} kb_id = meta.get("kb_id", "") out.append( QueryResult( chunk_id=str(chunk_id), content=content or "", score=float(score) if score is not None else 0.0, metadata=meta, document_id=str(document_id) if document_id is not None else "", kb_id=str(kb_id) if kb_id is not None else "", ) ) return out async def _retrieve_blend( self, query: str, kb_ids: list[str], top_k: int, ) -> list[QueryResult]: """双索引合并 — 语义 + 全文结果去重排序。 并行执行两种检索（各取 top_k），按 chunk_id 去重，分数归一化后加权融合。 若任一检索失败（如 embed_model 未配置），降级为另一种。 """ # 并行执行两种检索 embed_task = self._safe_retrieve_embedding(query, kb_ids, top_k) kw_task = self._retrieve_keywords(query, kb_ids, top_k) embed_results, kw_results = await asyncio.gather( embed_task, kw_task, return_exceptions=False ) # 归一化分数 embed_scores = _normalize_scores([r.score for r in embed_results]) for r, s in zip(embed_results, embed_scores): r.score = s * _BLEND_WEIGHT_EMBEDDING kw_scores = _normalize_scores([r.score for r in kw_results]) for r, s in zip(kw_results, kw_scores): r.score = s * _BLEND_WEIGHT_KEYWORDS # 合并去重 — 同 chunk_id 取最高分 merged: dict[str, QueryResult] = {} for r in (*embed_results, *kw_results): existing = merged.get(r.chunk_id) if existing is None or r.score > existing.score: merged[r.chunk_id] = r # 按分数降序，取 top_k results = sorted(merged.values(), key=lambda x: x.score, reverse=True) return results[:top_k] async def _safe_retrieve_embedding( self, query: str, kb_ids: list[str], top_k: int, ) -> list[QueryResult]: """embedding 检索的容错包装 — 失败时返回空列表（降级为纯关键词）。""" if self._embed_model is None: return [] try: return await self._retrieve_embedding(query, kb_ids, top_k) except Exception as e: logger.warning("Embedding retrieval failed, falling back to keywords only: %s", e) return [] __all__ = ["RetrievalEngine"]