256 lines
9.0 KiB
Python
256 lines
9.0 KiB
Python
"""双索引检索引擎 — 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, Any
|
||
|
||
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, Any] = 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"]
|