AstrBot 4.5.1__py3-none-any.whl → 4.5.2__py3-none-any.whl

astrbot/core/knowledge_base/retrieval/manager.py

@@ -4,18 +4,17 @@
 """
 
 import time
-
 from dataclasses import dataclass
-from typing import List
 
+from astrbot import logger
+from astrbot.core.db.vec_db.base import Result
+from astrbot.core.db.vec_db.faiss_impl import FaissVecDB
 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
@@ -53,6 +52,7 @@ class RetrievalManager:
             sparse_retriever: 稀疏检索器
             rank_fusion: 结果融合器
             kb_db: 知识库数据库实例
+
         """
         self.sparse_retriever = sparse_retriever
         self.rank_fusion = rank_fusion
@@ -61,11 +61,11 @@ class RetrievalManager:
     async def retrieve(
         self,
         query: str,
-        kb_ids: List[str],
+        kb_ids: list[str],
         kb_id_helper_map: dict[str, KBHelper],
         top_k_fusion: int = 20,
         top_m_final: int = 5,
-    ) -> List[RetrievalResult]:
+    ) -> list[RetrievalResult]:
         """混合检索
 
         流程:
@@ -82,6 +82,7 @@ class RetrievalManager:
 
         Returns:
             List[RetrievalResult]: 检索结果列表
+
         """
         if not kb_ids:
             return []
@@ -114,7 +115,7 @@ class RetrievalManager:
         )
         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."
+            f"Dense retrieval across {len(kb_ids)} bases took {time_end - time_start:.2f}s and returned {len(dense_results)} results.",
         )
 
         # 2. 稀疏检索
@@ -126,7 +127,7 @@ class RetrievalManager:
         )
         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."
+            f"Sparse retrieval across {len(kb_ids)} bases took {time_end - time_start:.2f}s and returned {len(sparse_results)} results.",
         )
 
         # 3. 结果融合
@@ -138,7 +139,7 @@ class RetrievalManager:
         )
         time_end = time.time()
         logger.debug(
-            f"Rank fusion took {time_end - time_start:.2f}s and returned {len(fused_results)} results."
+            f"Rank fusion took {time_end - time_start:.2f}s and returned {len(fused_results)} results.",
         )
 
         # 4. 转换为 RetrievalResult (获取元数据)
@@ -159,7 +160,7 @@ class RetrievalManager:
                             "chunk_index": fr.chunk_index,
                             "char_count": len(fr.content),
                         },
-                    )
+                    ),
                 )
 
         # 5. Rerank
@@ -188,7 +189,7 @@ class RetrievalManager:
     async def _dense_retrieve(
         self,
         query: str,
-        kb_ids: List[str],
+        kb_ids: list[str],
         kb_options: dict,
     ):
         """稠密检索 (向量相似度)
@@ -202,6 +203,7 @@ class RetrievalManager:
 
         Returns:
             List[Result]: 检索结果列表
+
         """
         all_results: list[Result] = []
         for kb_id in kb_ids:
@@ -233,10 +235,10 @@ class RetrievalManager:
     async def _rerank(
         self,
         query: str,
-        results: List[RetrievalResult],
+        results: list[RetrievalResult],
         top_k: int,
         rerank_provider: RerankProvider,
-    ) -> List[RetrievalResult]:
+    ) -> list[RetrievalResult]:
         """Rerank 重排序
 
         Args:
@@ -246,6 +248,7 @@ class RetrievalManager:
 
         Returns:
             List[RetrievalResult]: 重排序后的结果列表
+
         """
         if not results:
             return []

astrbot/core/knowledge_base/retrieval/rank_fusion.py

@@ -37,6 +37,7 @@ class RankFusion:
         Args:
             kb_db: 知识库数据库实例
             k: RRF 参数,用于平滑排名
+
         """
         self.kb_db = kb_db
         self.k = k
@@ -59,6 +60,7 @@ class RankFusion:
 
         Returns:
             List[FusedResult]: 融合后的结果列表
+
         """
         # 1. 构建排名映射
         dense_ranks = {
@@ -101,7 +103,9 @@ class RankFusion:
 
         # 4. 排序
         sorted_ids = sorted(
-            rrf_scores.keys(), key=lambda cid: rrf_scores[cid], reverse=True
+            rrf_scores.keys(),
+            key=lambda cid: rrf_scores[cid],
+            reverse=True,
         )[:top_k]
 
         # 5. 构建融合结果
@@ -118,7 +122,7 @@ class RankFusion:
                         kb_id=sr.kb_id,
                         content=sr.content,
                         score=rrf_scores[identifier],
-                    )
+                    ),
                 )
             elif identifier in vec_doc_id_to_dense:
                 # 从向量检索获取信息,需要从数据库获取块的详细信息
@@ -132,7 +136,7 @@ class RankFusion:
                         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

@@ -3,13 +3,15 @@
 使用 BM25 算法进行基于关键词的文档检索
 """
 
-import jieba
-import os
 import json
+import os
 from dataclasses import dataclass
+
+import jieba
 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
+from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase
 
 
 @dataclass
@@ -37,6 +39,7 @@ class SparseRetriever:
 
         Args:
             kb_db: 知识库数据库实例
+
         """
         self.kb_db = kb_db
         self._index_cache = {}  # 缓存 BM25 索引
@@ -64,6 +67,7 @@ class SparseRetriever:
 
         Returns:
             List[SparseResult]: 检索结果列表
+
         """
         # 1. 获取所有相关块
         top_k_sparse = 0
@@ -73,7 +77,9 @@ class SparseRetriever:
             if not vec_db:
                 continue
             result = await vec_db.document_storage.get_documents(
-                metadata_filters={}, limit=None, offset=None
+                metadata_filters={},
+                limit=None,
+                offset=None,
             )
             chunk_mds = [json.loads(doc["metadata"]) for doc in result]
             result = [
@@ -122,7 +128,7 @@ class SparseRetriever:
                     kb_id=chunk["kb_id"],
                     content=chunk["text"],
                     score=float(score),
-                )
+                ),
             )
 
         results.sort(key=lambda x: x.score, reverse=True)

astrbot/core/log.py

@@ -1,5 +1,4 @@
-"""
-日志系统, 用于支持核心组件和插件的日志记录, 提供了日志订阅功能
+"""日志系统, 用于支持核心组件和插件的日志记录, 提供了日志订阅功能
 
 const:
     CACHED_SIZE: 日志缓存大小, 用于限制缓存的日志数量
@@ -21,14 +20,14 @@ function:
 4. 订阅者可以使用 register() 方法注册到 LogBroker, 订阅日志流
 """
 
-import logging
-import colorlog
 import asyncio
+import logging
 import os
 import sys
-from collections import deque
 from asyncio import Queue
-from typing import List
+from collections import deque
+
+import colorlog
 
 # 日志缓存大小
 CACHED_SIZE = 200
@@ -52,6 +51,7 @@ def is_plugin_path(pathname):
 
     Returns:
         bool: 如果路径来自插件目录，则返回 True，否则返回 False
+
     """
     if not pathname:
         return False
@@ -68,6 +68,7 @@ def get_short_level_name(level_name):
 
     Returns:
         str: 四个字母的日志级别缩写
+
     """
     level_map = {
         "DEBUG": "DBUG",
@@ -87,13 +88,14 @@ class LogBroker:
 
     def __init__(self):
         self.log_cache = deque(maxlen=CACHED_SIZE)  # 环形缓冲区, 保存最近的日志
-        self.subscribers: List[Queue] = []  # 订阅者列表
+        self.subscribers: list[Queue] = []  # 订阅者列表
 
     def register(self) -> Queue:
         """注册新的订阅者, 并给每个订阅者返回一个带有日志缓存的队列
 
         Returns:
             Queue: 订阅者的队列, 可用于接收日志消息
+
         """
         q = Queue(maxsize=CACHED_SIZE + 10)
         self.subscribers.append(q)
@@ -104,6 +106,7 @@ class LogBroker:
 
         Args:
             q (Queue): 需要取消订阅的队列
+
         """
         self.subscribers.remove(q)
 
@@ -113,6 +116,7 @@ class LogBroker:
         Args:
             log_entry (dict): 日志消息, 包含日志级别和日志内容.
                 example: {"level": "INFO", "data": "This is a log message.", "time": "2023-10-01 12:00:00"}
+
         """
         self.log_cache.append(log_entry)
         for q in self.subscribers:
@@ -138,6 +142,7 @@ class LogQueueHandler(logging.Handler):
 
         Args:
             record (logging.LogRecord): 日志记录对象, 包含日志信息
+
         """
         log_entry = self.format(record)
         self.log_broker.publish(
@@ -145,7 +150,7 @@ class LogQueueHandler(logging.Handler):
                 "level": record.levelname,
                 "time": record.asctime,
                 "data": log_entry,
-            }
+            },
         )
 
 
@@ -164,6 +169,7 @@ class LogManager:
 
         Returns:
             logging.Logger: 返回配置好的日志记录器
+
         """
         logger = logging.getLogger(log_name)
         # 检查该logger或父级logger是否已经有处理器, 如果已经有处理器, 直接返回该logger, 避免重复配置
@@ -171,10 +177,10 @@ class LogManager:
             return logger
         # 如果logger没有处理器
         console_handler = logging.StreamHandler(
-            sys.stdout
+            sys.stdout,
         )  # 创建一个StreamHandler用于控制台输出
         console_handler.setLevel(
-            logging.DEBUG
+            logging.DEBUG,
         )  # 将日志级别设置为DEBUG(最低级别, 显示所有日志), *如果插件没有设置级别, 默认为DEBUG
 
         # 创建彩色日志格式化器, 输出日志格式为: [时间] [插件标签] [日志级别] [文件名:行号]: 日志消息
@@ -195,7 +201,8 @@ class LogManager:
 
         class FileNameFilter(logging.Filter):
             """文件名过滤器类, 用于修改日志记录的文件名格式
-            例如: 将文件路径 /path/to/file.py 转换为 file.<file> 格式"""
+            例如: 将文件路径 /path/to/file.py 转换为 file.<file> 格式
+            """
 
             # 获取这个文件和父文件夹的名字：<folder>.<file> 并且去除 .py
             def filter(self, record):
@@ -231,6 +238,7 @@ class LogManager:
         Args:
             logger (logging.Logger): 日志记录器
             log_broker (LogBroker): 日志代理类, 用于缓存和分发日志消息
+
         """
         handler = LogQueueHandler(log_broker)
         handler.setLevel(logging.DEBUG)
@@ -240,7 +248,7 @@ class LogManager:
             # 为队列处理器设置相同格式的formatter
             handler.setFormatter(
                 logging.Formatter(
-                    "[%(asctime)s] [%(short_levelname)s] %(plugin_tag)s[%(filename)s:%(lineno)d]: %(message)s"
-                )
+                    "[%(asctime)s] [%(short_levelname)s] %(plugin_tag)s[%(filename)s:%(lineno)d]: %(message)s",
+                ),
             )
         logger.addHandler(handler)