AstrBot 4.3.3__py3-none-any.whl → 4.5.0__py3-none-any.whl

astrbot/core/knowledge_base/retrieval/manager.py

@@ -0,0 +1,273 @@
+"""检索管理器
+
+协调稠密检索、稀疏检索和 Rerank,提供统一的检索接口
+"""
+
+import time
+
+from dataclasses import dataclass
+from typing import List
+
+from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase
+from astrbot.core.knowledge_base.retrieval.rank_fusion import RankFusion
+from astrbot.core.knowledge_base.retrieval.sparse_retriever import SparseRetriever
+from astrbot.core.provider.provider import RerankProvider
+from astrbot.core.db.vec_db.base import Result
+from astrbot.core.db.vec_db.faiss_impl import FaissVecDB
+from ..kb_helper import KBHelper
+from astrbot import logger
+
+
+@dataclass
+class RetrievalResult:
+    """检索结果"""
+
+    chunk_id: str
+    doc_id: str
+    doc_name: str
+    kb_id: str
+    kb_name: str
+    content: str
+    score: float
+    metadata: dict
+
+
+class RetrievalManager:
+    """检索管理器
+
+    职责:
+    - 协调稠密检索、稀疏检索和 Rerank
+    - 结果融合和排序
+    """
+
+    def __init__(
+        self,
+        sparse_retriever: SparseRetriever,
+        rank_fusion: RankFusion,
+        kb_db: KBSQLiteDatabase,
+    ):
+        """初始化检索管理器
+
+        Args:
+            vec_db_factory: 向量数据库工厂
+            sparse_retriever: 稀疏检索器
+            rank_fusion: 结果融合器
+            kb_db: 知识库数据库实例
+        """
+        self.sparse_retriever = sparse_retriever
+        self.rank_fusion = rank_fusion
+        self.kb_db = kb_db
+
+    async def retrieve(
+        self,
+        query: str,
+        kb_ids: List[str],
+        kb_id_helper_map: dict[str, KBHelper],
+        top_k_fusion: int = 20,
+        top_m_final: int = 5,
+    ) -> List[RetrievalResult]:
+        """混合检索
+
+        流程:
+        1. 稠密检索 (向量相似度)
+        2. 稀疏检索 (BM25)
+        3. 结果融合 (RRF)
+        4. Rerank 重排序
+
+        Args:
+            query: 查询文本
+            kb_ids: 知识库 ID 列表
+            top_m_final: 最终返回数量
+            enable_rerank: 是否启用 Rerank
+
+        Returns:
+            List[RetrievalResult]: 检索结果列表
+        """
+        if not kb_ids:
+            return []
+
+        kb_options: dict = {}
+        new_kb_ids = []
+        for kb_id in kb_ids:
+            kb_helper = kb_id_helper_map.get(kb_id)
+            if kb_helper:
+                kb = kb_helper.kb
+                kb_options[kb_id] = {
+                    "top_k_dense": kb.top_k_dense or 50,
+                    "top_k_sparse": kb.top_k_sparse or 50,
+                    "top_m_final": kb.top_m_final or 5,
+                    "vec_db": kb_helper.vec_db,
+                    "rerank_provider_id": kb.rerank_provider_id,
+                }
+                new_kb_ids.append(kb_id)
+            else:
+                logger.warning(f"知识库 ID {kb_id} 实例未找到, 已跳过该知识库的检索")
+
+        kb_ids = new_kb_ids
+
+        # 1. 稠密检索
+        time_start = time.time()
+        dense_results = await self._dense_retrieve(
+            query=query,
+            kb_ids=kb_ids,
+            kb_options=kb_options,
+        )
+        time_end = time.time()
+        logger.debug(
+            f"Dense retrieval across {len(kb_ids)} bases took {time_end - time_start:.2f}s and returned {len(dense_results)} results."
+        )
+
+        # 2. 稀疏检索
+        time_start = time.time()
+        sparse_results = await self.sparse_retriever.retrieve(
+            query=query,
+            kb_ids=kb_ids,
+            kb_options=kb_options,
+        )
+        time_end = time.time()
+        logger.debug(
+            f"Sparse retrieval across {len(kb_ids)} bases took {time_end - time_start:.2f}s and returned {len(sparse_results)} results."
+        )
+
+        # 3. 结果融合
+        time_start = time.time()
+        fused_results = await self.rank_fusion.fuse(
+            dense_results=dense_results,
+            sparse_results=sparse_results,
+            top_k=top_k_fusion,
+        )
+        time_end = time.time()
+        logger.debug(
+            f"Rank fusion took {time_end - time_start:.2f}s and returned {len(fused_results)} results."
+        )
+
+        # 4. 转换为 RetrievalResult (获取元数据)
+        retrieval_results = []
+        for fr in fused_results:
+            metadata_dict = await self.kb_db.get_document_with_metadata(fr.doc_id)
+            if metadata_dict:
+                retrieval_results.append(
+                    RetrievalResult(
+                        chunk_id=fr.chunk_id,
+                        doc_id=fr.doc_id,
+                        doc_name=metadata_dict["document"].doc_name,
+                        kb_id=fr.kb_id,
+                        kb_name=metadata_dict["knowledge_base"].kb_name,
+                        content=fr.content,
+                        score=fr.score,
+                        metadata={
+                            "chunk_index": fr.chunk_index,
+                            "char_count": len(fr.content),
+                        },
+                    )
+                )
+
+        # 5. Rerank
+        first_rerank = None
+        for kb_id in kb_ids:
+            vec_db: FaissVecDB = kb_options[kb_id]["vec_db"]
+            rerank_pi = kb_options[kb_id]["rerank_provider_id"]
+            if (
+                vec_db
+                and vec_db.rerank_provider
+                and rerank_pi
+                and rerank_pi == vec_db.rerank_provider.meta().id
+            ):
+                first_rerank = vec_db.rerank_provider
+                break
+        if first_rerank and retrieval_results:
+            retrieval_results = await self._rerank(
+                query=query,
+                results=retrieval_results,
+                top_k=top_m_final,
+                rerank_provider=first_rerank,
+            )
+
+        return retrieval_results[:top_m_final]
+
+    async def _dense_retrieve(
+        self,
+        query: str,
+        kb_ids: List[str],
+        kb_options: dict,
+    ):
+        """稠密检索 (向量相似度)
+
+        为每个知识库使用独立的向量数据库进行检索,然后合并结果。
+
+        Args:
+            query: 查询文本
+            kb_ids: 知识库 ID 列表
+            top_k: 返回结果数量
+
+        Returns:
+            List[Result]: 检索结果列表
+        """
+        all_results: list[Result] = []
+        for kb_id in kb_ids:
+            if kb_id not in kb_options:
+                continue
+            try:
+                vec_db: FaissVecDB = kb_options[kb_id]["vec_db"]
+                dense_k = int(kb_options[kb_id]["top_k_dense"])
+                vec_results = await vec_db.retrieve(
+                    query=query,
+                    k=dense_k,
+                    fetch_k=dense_k * 2,
+                    rerank=False,  # 稠密检索阶段不进行 rerank
+                    metadata_filters={"kb_id": kb_id},
+                )
+
+                all_results.extend(vec_results)
+            except Exception as e:
+                from astrbot.core import logger
+
+                logger.warning(f"知识库 {kb_id} 稠密检索失败: {e}")
+                continue
+
+        # 按相似度排序并返回 top_k
+        all_results.sort(key=lambda x: x.similarity, reverse=True)
+        # return all_results[: len(all_results) // len(kb_ids)]
+        return all_results
+
+    async def _rerank(
+        self,
+        query: str,
+        results: List[RetrievalResult],
+        top_k: int,
+        rerank_provider: RerankProvider,
+    ) -> List[RetrievalResult]:
+        """Rerank 重排序
+
+        Args:
+            query: 查询文本
+            results: 检索结果列表
+            top_k: 返回结果数量
+
+        Returns:
+            List[RetrievalResult]: 重排序后的结果列表
+        """
+        if not results:
+            return []
+
+        # 准备文档列表
+        docs = [r.content for r in results]
+
+        # 调用 Rerank Provider
+        rerank_results = await rerank_provider.rerank(
+            query=query,
+            documents=docs,
+        )
+
+        # 更新分数并重新排序
+        reranked_list = []
+        for rerank_result in rerank_results:
+            idx = rerank_result.index
+            if idx < len(results):
+                result = results[idx]
+                result.score = rerank_result.relevance_score
+                reranked_list.append(result)
+
+        reranked_list.sort(key=lambda x: x.score, reverse=True)
+
+        return reranked_list[:top_k]

astrbot/core/knowledge_base/retrieval/rank_fusion.py

@@ -0,0 +1,138 @@
+"""检索结果融合器
+
+使用 Reciprocal Rank Fusion (RRF) 算法融合稠密检索和稀疏检索的结果
+"""
+
+import json
+from dataclasses import dataclass
+
+from astrbot.core.db.vec_db.base import Result
+from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase
+from astrbot.core.knowledge_base.retrieval.sparse_retriever import SparseResult
+
+
+@dataclass
+class FusedResult:
+    """融合后的检索结果"""
+
+    chunk_id: str
+    chunk_index: int
+    doc_id: str
+    kb_id: str
+    content: str
+    score: float
+
+
+class RankFusion:
+    """检索结果融合器
+
+    职责:
+    - 融合稠密检索和稀疏检索的结果
+    - 使用 Reciprocal Rank Fusion (RRF) 算法
+    """
+
+    def __init__(self, kb_db: KBSQLiteDatabase, k: int = 60):
+        """初始化结果融合器
+
+        Args:
+            kb_db: 知识库数据库实例
+            k: RRF 参数,用于平滑排名
+        """
+        self.kb_db = kb_db
+        self.k = k
+
+    async def fuse(
+        self,
+        dense_results: list[Result],
+        sparse_results: list[SparseResult],
+        top_k: int = 20,
+    ) -> list[FusedResult]:
+        """融合稠密和稀疏检索结果
+
+        RRF 公式:
+        score(doc) = sum(1 / (k + rank_i))
+
+        Args:
+            dense_results: 稠密检索结果
+            sparse_results: 稀疏检索结果
+            top_k: 返回结果数量
+
+        Returns:
+            List[FusedResult]: 融合后的结果列表
+        """
+        # 1. 构建排名映射
+        dense_ranks = {
+            r.data["doc_id"]: (idx + 1) for idx, r in enumerate(dense_results)
+        }  # 这里的 doc_id 实际上是 chunk_id
+        sparse_ranks = {r.chunk_id: (idx + 1) for idx, r in enumerate(sparse_results)}
+
+        # 2. 收集所有唯一的 ID
+        # 需要统一为 chunk_id
+        all_chunk_ids = set()
+        vec_doc_id_to_dense: dict[str, Result] = {}  # vec_doc_id -> Result
+        chunk_id_to_sparse: dict[str, SparseResult] = {}  # chunk_id -> SparseResult
+
+        # 处理稀疏检索结果
+        for r in sparse_results:
+            all_chunk_ids.add(r.chunk_id)
+            chunk_id_to_sparse[r.chunk_id] = r
+
+        # 处理稠密检索结果 (需要转换 vec_doc_id 到 chunk_id)
+        for r in dense_results:
+            vec_doc_id = r.data["doc_id"]
+            all_chunk_ids.add(vec_doc_id)
+            vec_doc_id_to_dense[vec_doc_id] = r
+
+        # 3. 计算 RRF 分数
+        rrf_scores: dict[str, float] = {}
+
+        for identifier in all_chunk_ids:
+            score = 0.0
+
+            # 来自稠密检索的贡献
+            if identifier in dense_ranks:
+                score += 1.0 / (self.k + dense_ranks[identifier])
+
+            # 来自稀疏检索的贡献
+            if identifier in sparse_ranks:
+                score += 1.0 / (self.k + sparse_ranks[identifier])
+
+            rrf_scores[identifier] = score
+
+        # 4. 排序
+        sorted_ids = sorted(
+            rrf_scores.keys(), key=lambda cid: rrf_scores[cid], reverse=True
+        )[:top_k]
+
+        # 5. 构建融合结果
+        fused_results = []
+        for identifier in sorted_ids:
+            # 优先从稀疏检索获取完整信息
+            if identifier in chunk_id_to_sparse:
+                sr = chunk_id_to_sparse[identifier]
+                fused_results.append(
+                    FusedResult(
+                        chunk_id=sr.chunk_id,
+                        chunk_index=sr.chunk_index,
+                        doc_id=sr.doc_id,
+                        kb_id=sr.kb_id,
+                        content=sr.content,
+                        score=rrf_scores[identifier],
+                    )
+                )
+            elif identifier in vec_doc_id_to_dense:
+                # 从向量检索获取信息,需要从数据库获取块的详细信息
+                vec_result = vec_doc_id_to_dense[identifier]
+                chunk_md = json.loads(vec_result.data["metadata"])
+                fused_results.append(
+                    FusedResult(
+                        chunk_id=identifier,
+                        chunk_index=chunk_md["chunk_index"],
+                        doc_id=chunk_md["kb_doc_id"],
+                        kb_id=chunk_md["kb_id"],
+                        content=vec_result.data["text"],
+                        score=rrf_scores[identifier],
+                    )
+                )
+
+        return fused_results

astrbot/core/knowledge_base/retrieval/sparse_retriever.py

@@ -0,0 +1,130 @@
+"""稀疏检索器
+
+使用 BM25 算法进行基于关键词的文档检索
+"""
+
+import jieba
+import os
+import json
+from dataclasses import dataclass
+from rank_bm25 import BM25Okapi
+from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase
+from astrbot.core.db.vec_db.faiss_impl import FaissVecDB
+
+
+@dataclass
+class SparseResult:
+    """稀疏检索结果"""
+
+    chunk_index: int
+    chunk_id: str
+    doc_id: str
+    kb_id: str
+    content: str
+    score: float
+
+
+class SparseRetriever:
+    """BM25 稀疏检索器
+
+    职责:
+    - 基于关键词的文档检索
+    - 使用 BM25 算法计算相关度
+    """
+
+    def __init__(self, kb_db: KBSQLiteDatabase):
+        """初始化稀疏检索器
+
+        Args:
+            kb_db: 知识库数据库实例
+        """
+        self.kb_db = kb_db
+        self._index_cache = {}  # 缓存 BM25 索引
+
+        with open(
+            os.path.join(os.path.dirname(__file__), "hit_stopwords.txt"),
+            encoding="utf-8",
+        ) as f:
+            self.hit_stopwords = {
+                word.strip() for word in set(f.read().splitlines()) if word.strip()
+            }
+
+    async def retrieve(
+        self,
+        query: str,
+        kb_ids: list[str],
+        kb_options: dict,
+    ) -> list[SparseResult]:
+        """执行稀疏检索
+
+        Args:
+            query: 查询文本
+            kb_ids: 知识库 ID 列表
+            kb_options: 每个知识库的检索选项
+
+        Returns:
+            List[SparseResult]: 检索结果列表
+        """
+        # 1. 获取所有相关块
+        top_k_sparse = 0
+        chunks = []
+        for kb_id in kb_ids:
+            vec_db: FaissVecDB = kb_options.get(kb_id, {}).get("vec_db")
+            if not vec_db:
+                continue
+            result = await vec_db.document_storage.get_documents(
+                metadata_filters={}, limit=None, offset=None
+            )
+            chunk_mds = [json.loads(doc["metadata"]) for doc in result]
+            result = [
+                {
+                    "chunk_id": doc["doc_id"],
+                    "chunk_index": chunk_md["chunk_index"],
+                    "doc_id": chunk_md["kb_doc_id"],
+                    "kb_id": kb_id,
+                    "text": doc["text"],
+                }
+                for doc, chunk_md in zip(result, chunk_mds)
+            ]
+            chunks.extend(result)
+            top_k_sparse += kb_options.get(kb_id, {}).get("top_k_sparse", 50)
+
+        if not chunks:
+            return []
+
+        # 2. 准备文档和索引
+        corpus = [chunk["text"] for chunk in chunks]
+        tokenized_corpus = [list(jieba.cut(doc)) for doc in corpus]
+        tokenized_corpus = [
+            [word for word in doc if word not in self.hit_stopwords]
+            for doc in tokenized_corpus
+        ]
+
+        # 3. 构建 BM25 索引
+        bm25 = BM25Okapi(tokenized_corpus)
+
+        # 4. 执行检索
+        tokenized_query = list(jieba.cut(query))
+        tokenized_query = [
+            word for word in tokenized_query if word not in self.hit_stopwords
+        ]
+        scores = bm25.get_scores(tokenized_query)
+
+        # 5. 排序并返回 Top-K
+        results = []
+        for idx, score in enumerate(scores):
+            chunk = chunks[idx]
+            results.append(
+                SparseResult(
+                    chunk_id=chunk["chunk_id"],
+                    chunk_index=chunk["chunk_index"],
+                    doc_id=chunk["doc_id"],
+                    kb_id=chunk["kb_id"],
+                    content=chunk["text"],
+                    score=float(score),
+                )
+            )
+
+        results.sort(key=lambda x: x.score, reverse=True)
+        # return results[: len(results) // len(kb_ids)]
+        return results[:top_k_sparse]