diff --git a/.github/workflows/dashboard_ci.yml b/.github/workflows/dashboard_ci.yml index 9299dd6f..20562784 100644 --- a/.github/workflows/dashboard_ci.yml +++ b/.github/workflows/dashboard_ci.yml @@ -13,11 +13,18 @@ jobs: - name: Checkout repository uses: actions/checkout@v5 + - name: Setup Node.js + uses: actions/setup-node@v6 + with: + node-version: 'latest' + - name: npm install, build run: | cd dashboard - npm install - npm run build + npm install pnpm -g + pnpm install + pnpm i --save-dev @types/markdown-it + pnpm run build - name: Inject Commit SHA id: get_sha diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e0a0e314..8611e269 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -6,8 +6,20 @@ ci: autoupdate_schedule: weekly autoupdate_commit_msg: ":balloon: pre-commit autoupdate" repos: - - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.11.2 - hooks: - - id: ruff - - id: ruff-format +- repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.14.1 + hooks: + # Run the linter. + - id: ruff-check + types_or: [ python, pyi ] + args: [ --fix ] + # Run the formatter. + - id: ruff-format + types_or: [ python, pyi ] + +- repo: https://github.com/asottile/pyupgrade + rev: v3.21.0 + hooks: + - id: pyupgrade + args: [--py310-plus] diff --git a/README.md b/README.md index b65db314..6cffe774 100644 --- a/README.md +++ b/README.md @@ -4,14 +4,25 @@
-Soulter%2FAstrBot | Trendshift +
-[![GitHub release (latest by date)](https://img.shields.io/github/v/release/Soulter/AstrBot?style=for-the-badge&color=76bad9)](https://github.com/Soulter/AstrBot/releases/latest) +
+Soulter%2FAstrBot | Trendshift +Featured｜HelloGitHub +
+ +
+ +
+ python Docker pull -QQ_community -Telegram_community -![Dynamic JSON Badge](https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.soulter.top%2Fastrbot%2Fplugin-num&query=%24.result&suffix=%E4%B8%AA&style=for-the-badge&label=%E6%8F%92%E4%BB%B6%E5%B8%82%E5%9C%BA&cacheSeconds=3600) +QQ_community +Telegram_community + +
+ +
English ｜ 日本語 ｜ @@ -118,7 +129,7 @@ uv run main.py | QQ(OneBot) | ✔ | | Telegram | ✔ | | 企微应用 | ✔ | -| 企智机器人 | ✔ | +| 企微智能机器人 | ✔ | | 微信客服 | ✔ | | 微信公众号 | ✔ | | 飞书 | ✔ | diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 14013c09..33298e89 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -209,9 +209,38 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): ) continue + valid_params = {} # 参数过滤：只传递函数实际需要的参数 + + # 获取实际的 handler 函数 + if func_tool.handler: + logger.debug( + f"工具 {func_tool_name} 期望的参数: {func_tool.parameters}" + ) + if func_tool.parameters and func_tool.parameters.get("properties"): + expected_params = set(func_tool.parameters["properties"].keys()) + + valid_params = { + k: v + for k, v in func_tool_args.items() + if k in expected_params + } + + # 记录被忽略的参数 + ignored_params = set(func_tool_args.keys()) - set( + valid_params.keys() + ) + if ignored_params: + logger.warning( + f"工具 {func_tool_name} 忽略非期望参数: {ignored_params}" + ) + else: + # 如果没有 handler（如 MCP 工具），使用所有参数 + valid_params = func_tool_args + logger.warning(f"工具 {func_tool_name} 没有 handler，使用所有参数") + try: await self.agent_hooks.on_tool_start( - self.run_context, func_tool, func_tool_args + self.run_context, func_tool, valid_params ) except Exception as e: logger.error(f"Error in on_tool_start hook: {e}", exc_info=True) @@ -219,7 +248,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): executor = self.tool_executor.execute( tool=func_tool, run_context=self.run_context, - **func_tool_args, + **valid_params, # 只传递有效的参数 ) _final_resp: CallToolResult | None = None diff --git a/astrbot/core/astrbot_config_mgr.py b/astrbot/core/astrbot_config_mgr.py index 056135d0..0ee3f4fe 100644 --- a/astrbot/core/astrbot_config_mgr.py +++ b/astrbot/core/astrbot_config_mgr.py @@ -5,6 +5,7 @@ from astrbot.core.utils.shared_preferences import SharedPreferences from astrbot.core.config.astrbot_config import ASTRBOT_CONFIG_PATH from astrbot.core.config.default import DEFAULT_CONFIG from astrbot.core.platform.message_session import MessageSession +from astrbot.core.umop_config_router import UmopConfigRouter from astrbot.core.utils.astrbot_path import get_astrbot_config_path from typing import TypeVar, TypedDict @@ -15,14 +16,12 @@ class ConfInfo(TypedDict): """Configuration information for a specific session or platform.""" id: str # UUID of the configuration or "default" - umop: list[str] # Unified Message Origin Pattern name: str path: str # File name to the configuration file DEFAULT_CONFIG_CONF_INFO = ConfInfo( id="default", - umop=["::"], name="default", path=ASTRBOT_CONFIG_PATH, ) @@ -31,8 +30,14 @@ DEFAULT_CONFIG_CONF_INFO = ConfInfo( class AstrBotConfigManager: """A class to manage the system configuration of AstrBot, aka ACM""" - def __init__(self, default_config: AstrBotConfig, sp: SharedPreferences): + def __init__( + self, + default_config: AstrBotConfig, + ucr: UmopConfigRouter, + sp: SharedPreferences, + ): self.sp = sp + self.ucr = ucr self.confs: dict[str, AstrBotConfig] = {} """uuid / "default" -> AstrBotConfig""" self.confs["default"] = default_config @@ -63,24 +68,15 @@ class AstrBotConfigManager: ) continue - def _is_umo_match(self, p1: str, p2: str) -> bool: - """判断 p2 umo 是否逻辑包含于 p1 umo""" - p1_ls = p1.split(":") - p2_ls = p2.split(":") - - if len(p1_ls) != 3 or len(p2_ls) != 3: - return False # 非法格式 - - return all(p == "" or p == "*" or p == t for p, t in zip(p1_ls, p2_ls)) - def _load_conf_mapping(self, umo: str | MessageSession) -> ConfInfo: """获取指定 umo 的配置文件 uuid, 如果不存在则返回默认配置(返回 "default") Returns: ConfInfo: 包含配置文件的 uuid, 路径和名称等信息, 是一个 dict 类型 """ - # uuid -> { "umop": list, "path": str, "name": str } + # uuid -> { "path": str, "name": str } abconf_data = self._get_abconf_data() + if isinstance(umo, MessageSession): umo = str(umo) else: @@ -89,10 +85,13 @@ class AstrBotConfigManager: except Exception: return DEFAULT_CONFIG_CONF_INFO - for uuid_, meta in abconf_data.items(): - for pattern in meta["umop"]: - if self._is_umo_match(pattern, umo): - return ConfInfo(**meta, id=uuid_) + conf_id = self.ucr.get_conf_id_for_umop(umo) + if conf_id: + meta = abconf_data.get(conf_id) + if meta and isinstance(meta, dict): + # the bind relation between umo and conf is defined in ucr now, so we remove "umop" here + meta.pop("umop", None) + return ConfInfo(**meta, id=conf_id) return DEFAULT_CONFIG_CONF_INFO @@ -100,23 +99,14 @@ class AstrBotConfigManager: self, abconf_path: str, abconf_id: str, - umo_parts: list[str] | list[MessageSession], abconf_name: str | None = None, ) -> None: """保存配置文件的映射关系""" - for part in umo_parts: - if isinstance(part, MessageSession): - part = str(part) - elif not isinstance(part, str): - raise ValueError( - "umo_parts must be a list of strings or MessageSession instances" - ) abconf_data = self.sp.get( "abconf_mapping", {}, scope="global", scope_id="global" ) random_word = abconf_name or uuid.uuid4().hex[:8] abconf_data[abconf_id] = { - "umop": umo_parts, "path": abconf_path, "name": random_word, } @@ -153,29 +143,26 @@ class AstrBotConfigManager: def get_conf_list(self) -> list[ConfInfo]: """获取所有配置文件的元数据列表""" conf_list = [] - conf_list.append(DEFAULT_CONFIG_CONF_INFO) abconf_mapping = self._get_abconf_data() for uuid_, meta in abconf_mapping.items(): + if not isinstance(meta, dict): + continue + meta.pop("umop", None) conf_list.append(ConfInfo(**meta, id=uuid_)) + conf_list.append(DEFAULT_CONFIG_CONF_INFO) return conf_list def create_conf( self, - umo_parts: list[str] | list[MessageSession], config: dict = DEFAULT_CONFIG, name: str | None = None, ) -> str: - """ - umo 由三个部分组成 [platform_id]:[message_type]:[session_id]。 - - umo_parts 可以是 "::" (代表所有), 可以是 "[platform_id]::" (代表指定平台下的所有类型消息和会话)。 - """ conf_uuid = str(uuid.uuid4()) conf_file_name = f"abconf_{conf_uuid}.json" conf_path = os.path.join(get_astrbot_config_path(), conf_file_name) conf = AstrBotConfig(config_path=conf_path, default_config=config) conf.save_config() - self._save_conf_mapping(conf_file_name, conf_uuid, umo_parts, abconf_name=name) + self._save_conf_mapping(conf_file_name, conf_uuid, abconf_name=name) self.confs[conf_uuid] = conf return conf_uuid @@ -228,15 +215,12 @@ class AstrBotConfigManager: logger.info(f"成功删除配置文件 {conf_id}") return True - def update_conf_info( - self, conf_id: str, name: str | None = None, umo_parts: list[str] | None = None - ) -> bool: + def update_conf_info(self, conf_id: str, name: str | None = None) -> bool: """更新配置文件信息 Args: conf_id: 配置文件的 UUID name: 新的配置文件名称 (可选) - umo_parts: 新的 UMO 部分列表 (可选) Returns: bool: 更新是否成功 @@ -255,18 +239,6 @@ class AstrBotConfigManager: if name is not None: abconf_data[conf_id]["name"] = name - # 更新 UMO 部分 - if umo_parts is not None: - # 验证 UMO 部分格式 - for part in umo_parts: - if isinstance(part, MessageSession): - part = str(part) - elif not isinstance(part, str): - raise ValueError( - "umo_parts must be a list of strings or MessageSession instances" - ) - abconf_data[conf_id]["umop"] = umo_parts - # 保存更新 self.sp.put("abconf_mapping", abconf_data, scope="global", scope_id="global") self.abconf_data = abconf_data diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index ef063ae6..c0d16216 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -134,27 +134,11 @@ DEFAULT_CONFIG = { "persona": [], # deprecated "timezone": "Asia/Shanghai", "callback_api_base": "", - "default_kb_collection": "", # 默认知识库名称 + "default_kb_collection": "", # 默认知识库名称, 已经过时 "plugin_set": ["*"], # "*" 表示使用所有可用的插件, 空列表表示不使用任何插件 - "knowledge_base": { - "enabled": False, # 默认禁用，用户需要主动启用 - "embedding_provider_id": "", # 嵌入模型提供商 ID (为空时自动选择第一个) - "rerank_provider_id": "", # 重排序模型提供商 ID (为空时自动选择第一个) - "storage": { - "files_path": "data/knowledge_base", # 文件存储路径 - "vector_db_path": "data/knowledge_base/vectors", # 向量数据库路径 - }, - "chunking": { - "chunk_size": 512, # 文档块大小（字符数） - "chunk_overlap": 50, # 文档块重叠大小（字符数） - }, - "retrieval": { - "top_k_dense": 50, # 密集检索返回结果数 - "top_k_sparse": 50, # 稀疏检索返回结果数 - "top_m_final": 5, # 最终融合后返回的结果数 - "enable_rerank": True, # 是否启用重排序 - }, - }, + "kb_names": [], # 默认知识库名称列表 + "kb_fusion_top_k": 20, # 知识库检索融合阶段返回结果数量 + "kb_final_top_k": 5, # 知识库检索最终返回结果数量 } @@ -181,10 +165,11 @@ CONFIG_METADATA_2 = { "enable": False, "appid": "", "secret": "", + "is_sandbox": False, "callback_server_host": "0.0.0.0", "port": 6196, }, - "QQ 个人号(aiocqhttp)": { + "QQ 个人号(OneBot v11)": { "id": "default", "type": "aiocqhttp", "enable": False, @@ -192,7 +177,7 @@ CONFIG_METADATA_2 = { "ws_reverse_port": 6199, "ws_reverse_token": "", }, - "微信个人号(WeChatPadPro)": { + "WeChatPadPro": { "id": "wechatpadpro", "type": "wechatpadpro", "enable": False, @@ -287,6 +272,14 @@ CONFIG_METADATA_2 = { "misskey_default_visibility": "public", "misskey_local_only": False, "misskey_enable_chat": True, + # download / security options + "misskey_allow_insecure_downloads": False, + "misskey_download_timeout": 15, + "misskey_download_chunk_size": 65536, + "misskey_max_download_bytes": None, + "misskey_enable_file_upload": True, + "misskey_upload_concurrency": 3, + "misskey_upload_folder": "", }, "Slack": { "id": "slack", @@ -311,8 +304,26 @@ CONFIG_METADATA_2 = { "satori_heartbeat_interval": 10, "satori_reconnect_delay": 5, }, + # "WebChat": { + # "id": "webchat", + # "type": "webchat", + # "enable": False, + # "webchat_link_path": "", + # "webchat_present_type": "fullscreen", + # }, }, "items": { + # "webchat_link_path": { + # "description": "链接路径", + # "_special": "webchat_link_path", + # "type": "string", + # }, + # "webchat_present_type": { + # "_special": "webchat_present_type", + # "description": "展现形式", + # "type": "string", + # "options": ["fullscreen", "embedded"], + # }, "satori_api_base_url": { "description": "Satori API 终结点", "type": "string", @@ -415,6 +426,41 @@ CONFIG_METADATA_2 = { "type": "bool", "hint": "启用后，机器人将会监听和响应私信聊天消息", }, + "misskey_enable_file_upload": { + "description": "启用文件上传到 Misskey", + "type": "bool", + "hint": "启用后，适配器会尝试将消息链中的文件上传到 Misskey。URL 文件会先尝试服务器端上传，异步上传失败时会回退到下载后本地上传。", + }, + "misskey_allow_insecure_downloads": { + "description": "允许不安全下载（禁用 SSL 验证）", + "type": "bool", + "hint": "当远端服务器存在证书问题导致无法正常下载时，自动禁用 SSL 验证作为回退方案。适用于某些图床的证书配置问题。启用有安全风险，仅在必要时使用。", + }, + "misskey_download_timeout": { + "description": "远端下载超时时间（秒）", + "type": "int", + "hint": "下载远程文件时的超时时间（秒），用于异步上传回退到本地上传的场景。", + }, + "misskey_download_chunk_size": { + "description": "流式下载分块大小（字节）", + "type": "int", + "hint": "流式下载和计算 MD5 时使用的每次读取字节数，过小会增加开销，过大会占用内存。", + }, + "misskey_max_download_bytes": { + "description": "最大允许下载字节数（超出则中止）", + "type": "int", + "hint": "如果希望限制下载文件的最大大小以防止 OOM，请填写最大字节数；留空或 null 表示不限制。", + }, + "misskey_upload_concurrency": { + "description": "并发上传限制", + "type": "int", + "hint": "同时进行的文件上传任务上限（整数，默认 3）。", + }, + "misskey_upload_folder": { + "description": "上传到网盘的目标文件夹 ID", + "type": "string", + "hint": "可选：填写 Misskey 网盘中目标文件夹的 ID，上传的文件将放置到该文件夹内。留空则使用账号网盘根目录。", + }, "telegram_command_register": { "description": "Telegram 命令注册", "type": "bool", @@ -466,19 +512,18 @@ CONFIG_METADATA_2 = { "hint": "启用后，机器人可以接收到频道的私聊消息。", }, "ws_reverse_host": { - "description": "反向 Websocket 主机地址(AstrBot 为服务器端)", + "description": "反向 Websocket 主机", "type": "string", - "hint": "aiocqhttp 适配器的反向 Websocket 服务器 IP 地址，不包含端口号。", + "hint": "AstrBot 将作为服务器端。", }, "ws_reverse_port": { "description": "反向 Websocket 端口", "type": "int", - "hint": "aiocqhttp 适配器的反向 Websocket 端口。", }, "ws_reverse_token": { "description": "反向 Websocket Token", "type": "string", - "hint": "aiocqhttp 适配器的反向 Websocket Token。未设置则不启用 Token 验证。", + "hint": "反向 Websocket Token。未设置则不启用 Token 验证。", }, "wecom_ai_bot_name": { "description": "企业微信智能机器人的名字", @@ -2019,6 +2064,9 @@ CONFIG_METADATA_2 = { "default_kb_collection": { "type": "string", }, + "kb_names": {"type": "list", "items": {"type": "string"}}, + "kb_fusion_top_k": {"type": "int", "default": 20}, + "kb_final_top_k": {"type": "int", "default": 5}, }, }, } @@ -2097,10 +2145,22 @@ CONFIG_METADATA_3 = { "description": "知识库", "type": "object", "items": { - "default_kb_collection": { - "description": "默认使用的知识库", - "type": "string", + "kb_names": { + "description": "知识库列表", + "type": "list", + "items": {"type": "string"}, "_special": "select_knowledgebase", + "hint": "支持多选", + }, + "kb_fusion_top_k": { + "description": "融合检索结果数", + "type": "int", + "hint": "多个知识库检索结果融合后的返回结果数量", + }, + "kb_final_top_k": { + "description": "最终返回结果数", + "type": "int", + "hint": "从知识库中检索到的结果数量，越大可能获得越多相关信息，但也可能引入噪音。建议根据实际需求调整", }, }, }, @@ -2194,7 +2254,7 @@ CONFIG_METADATA_3 = { "provider_settings.wake_prefix": { "description": "LLM 聊天额外唤醒前缀 ", "type": "string", - "hint": "例子: 如果唤醒前缀为 `/`, 额外聊天唤醒前缀为 `chat`，则需要 `/chat` 才会触发 LLM 请求。默认为空。", + "hint": "如果唤醒前缀为 `/`, 额外聊天唤醒前缀为 `chat`，则需要 `/chat` 才会触发 LLM 请求。默认为空。", }, "provider_settings.prompt_prefix": { "description": "用户提示词", diff --git a/astrbot/core/core_lifecycle.py b/astrbot/core/core_lifecycle.py index cc0122a1..3d4b28c0 100644 --- a/astrbot/core/core_lifecycle.py +++ b/astrbot/core/core_lifecycle.py @@ -17,7 +17,6 @@ import os from .event_bus import EventBus from . import astrbot_config, html_renderer from asyncio import Queue -from typing import List from astrbot.core.pipeline.scheduler import PipelineScheduler, PipelineContext from astrbot.core.star import PluginManager from astrbot.core.platform.manager import PlatformManager @@ -26,15 +25,17 @@ from astrbot.core.persona_mgr import PersonaManager from astrbot.core.provider.manager import ProviderManager from astrbot.core import LogBroker from astrbot.core.db import BaseDatabase +from astrbot.core.db.migration.migra_45_to_46 import migrate_45_to_46 from astrbot.core.updator import AstrBotUpdator from astrbot.core import logger, sp from astrbot.core.config.default import VERSION from astrbot.core.conversation_mgr import ConversationManager from astrbot.core.platform_message_history_mgr import PlatformMessageHistoryManager +from astrbot.core.umop_config_router import UmopConfigRouter from astrbot.core.astrbot_config_mgr import AstrBotConfigManager from astrbot.core.star.star_handler import star_handlers_registry, EventType from astrbot.core.star.star_handler import star_map -from astrbot.core.knowledge_base.kb_manager_lifecycle import KnowledgeBaseManager +from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager class AstrBotCoreLifecycle: @@ -85,11 +86,21 @@ class AstrBotCoreLifecycle: await html_renderer.initialize() + # 初始化 UMOP 配置路由器 + self.umop_config_router = UmopConfigRouter(sp=sp) + # 初始化 AstrBot 配置管理器 self.astrbot_config_mgr = AstrBotConfigManager( - default_config=self.astrbot_config, sp=sp + default_config=self.astrbot_config, ucr=self.umop_config_router, sp=sp ) + # 4.5 to 4.6 migration for umop_config_router + try: + await migrate_45_to_46(self.astrbot_config_mgr, self.umop_config_router) + except Exception as e: + logger.error(f"Migration from version 4.5 to 4.6 failed: {e!s}") + logger.error(traceback.format_exc()) + # 初始化事件队列 self.event_queue = Queue() @@ -112,9 +123,7 @@ class AstrBotCoreLifecycle: self.platform_message_history_manager = PlatformMessageHistoryManager(self.db) # 初始化知识库管理器 - self.kb_manager = KnowledgeBaseManager( - self.astrbot_config, self.db, self.provider_manager - ) + self.kb_manager = KnowledgeBaseManager(self.provider_manager) # 初始化提供给插件的上下文 self.star_context = Context( @@ -141,10 +150,6 @@ class AstrBotCoreLifecycle: await self.kb_manager.initialize() - # 注册知识库会话生命周期钩子（零侵入级联清理） - if self.kb_manager.is_initialized: - self.kb_manager.register_session_lifecycle_hooks(self.conversation_manager) - # 初始化消息事件流水线调度器 self.pipeline_scheduler_mapping = await self.load_pipeline_scheduler() @@ -160,7 +165,7 @@ class AstrBotCoreLifecycle: self.start_time = int(time.time()) # 初始化当前任务列表 - self.curr_tasks: List[asyncio.Task] = [] + self.curr_tasks: list[asyncio.Task] = [] # 根据配置实例化各个平台适配器 await self.platform_manager.initialize() @@ -267,7 +272,7 @@ class AstrBotCoreLifecycle: target=self.astrbot_updator._reboot, name="restart", daemon=True ).start() - def load_platform(self) -> List[asyncio.Task]: + def load_platform(self) -> list[asyncio.Task]: """加载平台实例并返回所有平台实例的异步任务列表""" tasks = [] platform_insts = self.platform_manager.get_insts() diff --git a/astrbot/core/db/migration/migra_45_to_46.py b/astrbot/core/db/migration/migra_45_to_46.py new file mode 100644 index 00000000..8a1dc5de --- /dev/null +++ b/astrbot/core/db/migration/migra_45_to_46.py @@ -0,0 +1,44 @@ +from astrbot.api import logger, sp +from astrbot.core.astrbot_config_mgr import AstrBotConfigManager +from astrbot.core.umop_config_router import UmopConfigRouter + + +async def migrate_45_to_46(acm: AstrBotConfigManager, ucr: UmopConfigRouter): + abconf_data = acm.abconf_data + + if not isinstance(abconf_data, dict): + # should be unreachable + logger.warning( + f"migrate_45_to_46: abconf_data is not a dict (type={type(abconf_data)}). Value: {abconf_data!r}" + ) + return + + # 如果任何一项带有 umop，则说明需要迁移 + need_migration = False + for conf_id, conf_info in abconf_data.items(): + if isinstance(conf_info, dict) and "umop" in conf_info: + need_migration = True + break + + if not need_migration: + return + + logger.info("Starting migration from version 4.5 to 4.6") + + # extract umo->conf_id mapping + umo_to_conf_id = {} + for conf_id, conf_info in abconf_data.items(): + if isinstance(conf_info, dict) and "umop" in conf_info: + umop_ls = conf_info.pop("umop") + if not isinstance(umop_ls, list): + continue + for umo in umop_ls: + if isinstance(umo, str) and umo not in umo_to_conf_id: + umo_to_conf_id[umo] = conf_id + + # update the abconf data + await sp.global_put("abconf_mapping", abconf_data) + # update the umop config router + await ucr.update_routing_data(umo_to_conf_id) + + logger.info("Migration from version 45 to 46 completed successfully") diff --git a/astrbot/core/db/vec_db/base.py b/astrbot/core/db/vec_db/base.py index d6e5e56b..27fc9f3f 100644 --- a/astrbot/core/db/vec_db/base.py +++ b/astrbot/core/db/vec_db/base.py @@ -16,14 +16,42 @@ class BaseVecDB: pass @abc.abstractmethod - async def insert(self, content: str, metadata: dict = None, id: str = None) -> int: + async def insert( + self, content: str, metadata: dict | None = None, id: str | None = None + ) -> int: """ 插入一条文本和其对应向量，自动生成 ID 并保持一致性。 """ ... @abc.abstractmethod - async def retrieve(self, query: str, top_k: int = 5) -> list[Result]: + async def insert_batch( + self, + contents: list[str], + metadatas: list[dict] | None = None, + ids: list[str] | None = None, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> int: + """ + 批量插入文本和其对应向量，自动生成 ID 并保持一致性。 + + Args: + progress_callback: 进度回调函数，接收参数 (current, total) + """ + ... + + @abc.abstractmethod + async def retrieve( + self, + query: str, + top_k: int = 5, + fetch_k: int = 20, + rerank: bool = False, + metadata_filters: dict | None = None, + ) -> list[Result]: """ 搜索最相似的文档。 Args: @@ -44,3 +72,6 @@ class BaseVecDB: bool: 删除是否成功 """ ... + + @abc.abstractmethod + async def close(self): ... diff --git a/astrbot/core/db/vec_db/faiss_impl/document_storage.py b/astrbot/core/db/vec_db/faiss_impl/document_storage.py index ee44da66..1feeb9b9 100644 --- a/astrbot/core/db/vec_db/faiss_impl/document_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/document_storage.py @@ -1,59 +1,219 @@ -import aiosqlite import os +import json +from datetime import datetime +from contextlib import asynccontextmanager + +from sqlalchemy import Text, Column +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine +from sqlalchemy.orm import sessionmaker +from sqlmodel import Field, SQLModel, select, col, func, text, MetaData + + +class BaseDocModel(SQLModel, table=False): + metadata = MetaData() + + +class Document(BaseDocModel, table=True): + """SQLModel for documents table.""" + + __tablename__ = "documents" # type: ignore + + id: int | None = Field( + default=None, primary_key=True, sa_column_kwargs={"autoincrement": True} + ) + doc_id: str = Field(nullable=False) + text: str = Field(nullable=False) + metadata_: str | None = Field(default=None, sa_column=Column("metadata", Text)) + created_at: datetime | None = Field(default=None) + updated_at: datetime | None = Field(default=None) class DocumentStorage: def __init__(self, db_path: str): self.db_path = db_path - self.connection = None + self.DATABASE_URL = f"sqlite+aiosqlite:///{db_path}" + self.engine: AsyncEngine | None = None + self.async_session_maker: sessionmaker | None = None self.sqlite_init_path = os.path.join( os.path.dirname(__file__), "sqlite_init.sql" ) async def initialize(self): """Initialize the SQLite database and create the documents table if it doesn't exist.""" - if not os.path.exists(self.db_path): - await self.connect() - async with self.connection.cursor() as cursor: - with open(self.sqlite_init_path, "r", encoding="utf-8") as f: - sql_script = f.read() - await cursor.executescript(sql_script) - await self.connection.commit() - else: - await self.connect() + await self.connect() + async with self.engine.begin() as conn: # type: ignore + # Create tables using SQLModel + await conn.run_sync(BaseDocModel.metadata.create_all) + + try: + await conn.execute( + text( + "ALTER TABLE documents ADD COLUMN kb_doc_id TEXT " + "GENERATED ALWAYS AS (json_extract(metadata, '$.kb_doc_id')) STORED" + ) + ) + await conn.execute( + text( + "ALTER TABLE documents ADD COLUMN user_id TEXT " + "GENERATED ALWAYS AS (json_extract(metadata, '$.user_id')) STORED" + ) + ) + + # Create indexes + await conn.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_documents_kb_doc_id ON documents(kb_doc_id)" + ) + ) + await conn.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_documents_user_id ON documents(user_id)" + ) + ) + except BaseException: + pass + + await conn.commit() async def connect(self): """Connect to the SQLite database.""" - self.connection = await aiosqlite.connect(self.db_path) + if self.engine is None: + self.engine = create_async_engine( + self.DATABASE_URL, + echo=False, + future=True, + ) + self.async_session_maker = sessionmaker( + self.engine, # type: ignore + class_=AsyncSession, + expire_on_commit=False, + ) # type: ignore - async def get_documents(self, metadata_filters: dict, ids: list = None): + @asynccontextmanager + async def get_session(self): + """Context manager for database sessions.""" + async with self.async_session_maker() as session: # type: ignore + yield session + + async def get_documents( + self, + metadata_filters: dict, + ids: list | None = None, + offset: int | None = 0, + limit: int | None = 100, + ) -> list[dict]: """Retrieve documents by metadata filters and ids. Args: metadata_filters (dict): The metadata filters to apply. + ids (list | None): Optional list of document IDs to filter. + offset (int | None): Offset for pagination. + limit (int | None): Limit for pagination. Returns: - list: The list of document IDs(primary key, not doc_id) that match the filters. + list: The list of documents that match the filters. """ - # metadata filter -> SQL WHERE clause - where_clauses = [] - values = [] - for key, val in metadata_filters.items(): - where_clauses.append(f"json_extract(metadata, '$.{key}') = ?") - values.append(val) - if ids is not None and len(ids) > 0: - ids = [str(i) for i in ids if i != -1] - where_clauses.append("id IN ({})".format(",".join("?" * len(ids)))) - values.extend(ids) - where_sql = " AND ".join(where_clauses) or "1=1" + assert self.engine is not None, "Database connection is not initialized." - result = [] - async with self.connection.cursor() as cursor: - sql = "SELECT * FROM documents WHERE " + where_sql - await cursor.execute(sql, values) - for row in await cursor.fetchall(): - result.append(await self.tuple_to_dict(row)) - return result + async with self.get_session() as session: + query = select(Document) + + for key, val in metadata_filters.items(): + query = query.where( + text(f"json_extract(metadata, '$.{key}') = :filter_{key}") + ).params(**{f"filter_{key}": val}) + + if ids is not None and len(ids) > 0: + valid_ids = [int(i) for i in ids if i != -1] + if valid_ids: + query = query.where(col(Document.id).in_(valid_ids)) + + if limit is not None: + query = query.limit(limit) + if offset is not None: + query = query.offset(offset) + + result = await session.execute(query) + documents = result.scalars().all() + + return [self._document_to_dict(doc) for doc in documents] + + async def insert_document(self, doc_id: str, text: str, metadata: dict) -> int: + """Insert a single document and return its integer ID. + + Args: + doc_id (str): The document ID (UUID string). + text (str): The document text. + metadata (dict): The document metadata. + + Returns: + int: The integer ID of the inserted document. + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + async with session.begin(): + document = Document( + doc_id=doc_id, + text=text, + metadata_=json.dumps(metadata), + created_at=datetime.now(), + updated_at=datetime.now(), + ) + session.add(document) + await session.flush() # Flush to get the ID + return document.id # type: ignore + + async def insert_documents_batch( + self, doc_ids: list[str], texts: list[str], metadatas: list[dict] + ) -> list[int]: + """Batch insert documents and return their integer IDs. + + Args: + doc_ids (list[str]): List of document IDs (UUID strings). + texts (list[str]): List of document texts. + metadatas (list[dict]): List of document metadata. + + Returns: + list[int]: List of integer IDs of the inserted documents. + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + async with session.begin(): + import json + + documents = [] + for doc_id, text, metadata in zip(doc_ids, texts, metadatas): + document = Document( + doc_id=doc_id, + text=text, + metadata_=json.dumps(metadata), + created_at=datetime.now(), + updated_at=datetime.now(), + ) + documents.append(document) + session.add(document) + + await session.flush() # Flush to get all IDs + return [doc.id for doc in documents] # type: ignore + + async def delete_document_by_doc_id(self, doc_id: str): + """Delete a document by its doc_id. + + Args: + doc_id (str): The doc_id of the document to delete. + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + async with session.begin(): + query = select(Document).where(col(Document.doc_id) == doc_id) + result = await session.execute(query) + document = result.scalar_one_or_none() + + if document: + await session.delete(document) async def get_document_by_doc_id(self, doc_id: str): """Retrieve a document by its doc_id. @@ -62,28 +222,85 @@ class DocumentStorage: doc_id (str): The doc_id of the document to retrieve. Returns: - dict: The document data. + dict: The document data or None if not found. """ - async with self.connection.cursor() as cursor: - await cursor.execute("SELECT * FROM documents WHERE doc_id = ?", (doc_id,)) - row = await cursor.fetchone() - if row: - return await self.tuple_to_dict(row) - else: - return None + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + query = select(Document).where(col(Document.doc_id) == doc_id) + result = await session.execute(query) + document = result.scalar_one_or_none() + + if document: + return self._document_to_dict(document) + return None async def update_document_by_doc_id(self, doc_id: str, new_text: str): - """Retrieve a document by its doc_id. + """Update a document by its doc_id. Args: doc_id (str): The doc_id. new_text (str): The new text to update the document with. """ - async with self.connection.cursor() as cursor: - await cursor.execute( - "UPDATE documents SET text = ? WHERE doc_id = ?", (new_text, doc_id) - ) - await self.connection.commit() + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + async with session.begin(): + query = select(Document).where(col(Document.doc_id) == doc_id) + result = await session.execute(query) + document = result.scalar_one_or_none() + + if document: + document.text = new_text + document.updated_at = datetime.now() + session.add(document) + + async def delete_documents(self, metadata_filters: dict): + """Delete documents by their metadata filters. + + Args: + metadata_filters (dict): The metadata filters to apply. + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + async with session.begin(): + query = select(Document) + + for key, val in metadata_filters.items(): + query = query.where( + text(f"json_extract(metadata, '$.{key}') = :filter_{key}") + ).params(**{f"filter_{key}": val}) + + result = await session.execute(query) + documents = result.scalars().all() + + for doc in documents: + await session.delete(doc) + + async def count_documents(self, metadata_filters: dict | None = None) -> int: + """Count documents in the database. + + Args: + metadata_filters (dict | None): Metadata filters to apply. + + Returns: + int: The count of documents. + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + query = select(func.count(col(Document.id))) + + if metadata_filters: + for key, val in metadata_filters.items(): + query = query.where( + text(f"json_extract(metadata, '$.{key}') = :filter_{key}") + ).params(**{f"filter_{key}": val}) + + result = await session.execute(query) + count = result.scalar_one_or_none() + return count if count is not None else 0 async def get_user_ids(self) -> list[str]: """Retrieve all user IDs from the documents table. @@ -91,11 +308,38 @@ class DocumentStorage: Returns: list: A list of user IDs. """ - async with self.connection.cursor() as cursor: - await cursor.execute("SELECT DISTINCT user_id FROM documents") - rows = await cursor.fetchall() + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + query = text( + "SELECT DISTINCT user_id FROM documents WHERE user_id IS NOT NULL" + ) + result = await session.execute(query) + rows = result.fetchall() return [row[0] for row in rows] + def _document_to_dict(self, document: Document) -> dict: + """Convert a Document model to a dictionary. + + Args: + document (Document): The document to convert. + + Returns: + dict: The converted dictionary. + """ + return { + "id": document.id, + "doc_id": document.doc_id, + "text": document.text, + "metadata": document.metadata_, + "created_at": document.created_at.isoformat() + if isinstance(document.created_at, datetime) + else document.created_at, + "updated_at": document.updated_at.isoformat() + if isinstance(document.updated_at, datetime) + else document.updated_at, + } + async def tuple_to_dict(self, row): """Convert a tuple to a dictionary. @@ -104,6 +348,8 @@ class DocumentStorage: Returns: dict: The converted dictionary. + + Note: This method is kept for backward compatibility but is no longer used internally. """ return { "id": row[0], @@ -116,6 +362,7 @@ class DocumentStorage: async def close(self): """Close the connection to the SQLite database.""" - if self.connection: - await self.connection.close() - self.connection = None + if self.engine: + await self.engine.dispose() + self.engine = None + self.async_session_maker = None diff --git a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py index 18ee3189..2c0cc8df 100644 --- a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py @@ -9,7 +9,7 @@ import numpy as np class EmbeddingStorage: - def __init__(self, dimension: int, path: str = None): + def __init__(self, dimension: int, path: str | None = None): self.dimension = dimension self.path = path self.index = None @@ -18,7 +18,6 @@ class EmbeddingStorage: else: base_index = faiss.IndexFlatL2(dimension) self.index = faiss.IndexIDMap(base_index) - self.storage = {} async def insert(self, vector: np.ndarray, id: int): """插入向量 @@ -29,12 +28,29 @@ class EmbeddingStorage: Raises: ValueError: 如果向量的维度与存储的维度不匹配 """ + assert self.index is not None, "FAISS index is not initialized." if vector.shape[0] != self.dimension: raise ValueError( f"向量维度不匹配, 期望: {self.dimension}, 实际: {vector.shape[0]}" ) self.index.add_with_ids(vector.reshape(1, -1), np.array([id])) - self.storage[id] = vector + await self.save_index() + + async def insert_batch(self, vectors: np.ndarray, ids: list[int]): + """批量插入向量 + + Args: + vectors (np.ndarray): 要插入的向量数组 + ids (list[int]): 向量的ID列表 + Raises: + ValueError: 如果向量的维度与存储的维度不匹配 + """ + assert self.index is not None, "FAISS index is not initialized." + if vectors.shape[1] != self.dimension: + raise ValueError( + f"向量维度不匹配, 期望: {self.dimension}, 实际: {vectors.shape[1]}" + ) + self.index.add_with_ids(vectors, np.array(ids)) await self.save_index() async def search(self, vector: np.ndarray, k: int) -> tuple: @@ -46,10 +62,22 @@ class EmbeddingStorage: Returns: tuple: (距离, 索引) """ + assert self.index is not None, "FAISS index is not initialized." faiss.normalize_L2(vector) distances, indices = self.index.search(vector, k) return distances, indices + async def delete(self, ids: list[int]): + """删除向量 + + Args: + ids (list[int]): 要删除的向量ID列表 + """ + assert self.index is not None, "FAISS index is not initialized." + id_array = np.array(ids, dtype=np.int64) + self.index.remove_ids(id_array) + await self.save_index() + async def save_index(self): """保存索引 diff --git a/astrbot/core/db/vec_db/faiss_impl/vec_db.py b/astrbot/core/db/vec_db/faiss_impl/vec_db.py index 7c2ae1c0..8a21538e 100644 --- a/astrbot/core/db/vec_db/faiss_impl/vec_db.py +++ b/astrbot/core/db/vec_db/faiss_impl/vec_db.py @@ -1,11 +1,12 @@ import uuid -import json +import time import numpy as np from .document_storage import DocumentStorage from .embedding_storage import EmbeddingStorage from ..base import Result, BaseVecDB from astrbot.core.provider.provider import EmbeddingProvider from astrbot.core.provider.provider import RerankProvider +from astrbot import logger class FaissVecDB(BaseVecDB): @@ -44,18 +45,56 @@ class FaissVecDB(BaseVecDB): vector = await self.embedding_provider.get_embedding(content) vector = np.array(vector, dtype=np.float32) - async with self.document_storage.connection.cursor() as cursor: - await cursor.execute( - "INSERT INTO documents (doc_id, text, metadata) VALUES (?, ?, ?)", - (str_id, content, json.dumps(metadata)), - ) - await self.document_storage.connection.commit() - result = await self.document_storage.get_document_by_doc_id(str_id) - int_id = result["id"] - # 插入向量到 FAISS - await self.embedding_storage.insert(vector, int_id) - return int_id + # 使用 DocumentStorage 的方法插入文档 + int_id = await self.document_storage.insert_document(str_id, content, metadata) + + # 插入向量到 FAISS + await self.embedding_storage.insert(vector, int_id) + return int_id + + async def insert_batch( + self, + contents: list[str], + metadatas: list[dict] | None = None, + ids: list[str] | None = None, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> list[int]: + """ + 批量插入文本和其对应向量，自动生成 ID 并保持一致性。 + + Args: + progress_callback: 进度回调函数，接收参数 (current, total) + """ + metadatas = metadatas or [{} for _ in contents] + ids = ids or [str(uuid.uuid4()) for _ in contents] + + start = time.time() + logger.debug(f"Generating embeddings for {len(contents)} contents...") + vectors = await self.embedding_provider.get_embeddings_batch( + contents, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=progress_callback, + ) + end = time.time() + logger.debug( + f"Generated embeddings for {len(contents)} contents in {end - start:.2f} seconds." + ) + + # 使用 DocumentStorage 的批量插入方法 + int_ids = await self.document_storage.insert_documents_batch( + ids, contents, metadatas + ) + + # 批量插入向量到 FAISS + vectors_array = np.array(vectors).astype("float32") + await self.embedding_storage.insert_batch(vectors_array, int_ids) + return int_ids async def retrieve( self, @@ -119,23 +158,42 @@ class FaissVecDB(BaseVecDB): return top_k_results - async def delete(self, doc_id: int): + async def delete(self, doc_id: str): """ - 删除一条文档 + 删除一条文档块（chunk） """ - await self.document_storage.connection.execute( - "DELETE FROM documents WHERE doc_id = ?", (doc_id,) - ) - await self.document_storage.connection.commit() + # 获得对应的 int id + result = await self.document_storage.get_document_by_doc_id(doc_id) + int_id = result["id"] if result else None + if int_id is None: + return + + # 使用 DocumentStorage 的删除方法 + await self.document_storage.delete_document_by_doc_id(doc_id) + await self.embedding_storage.delete([int_id]) async def close(self): await self.document_storage.close() - async def count_documents(self) -> int: + async def count_documents(self, metadata_filter: dict | None = None) -> int: """ 计算文档数量 + + Args: + metadata_filter (dict | None): 元数据过滤器 """ - async with self.document_storage.connection.cursor() as cursor: - await cursor.execute("SELECT COUNT(*) FROM documents") - count = await cursor.fetchone() - return count[0] if count else 0 + count = await self.document_storage.count_documents( + metadata_filters=metadata_filter or {} + ) + return count + + async def delete_documents(self, metadata_filters: dict): + """ + 根据元数据过滤器删除文档 + """ + docs = await self.document_storage.get_documents( + metadata_filters=metadata_filters, offset=None, limit=None + ) + doc_ids: list[int] = [doc["id"] for doc in docs] + await self.embedding_storage.delete(doc_ids) + await self.document_storage.delete_documents(metadata_filters=metadata_filters) diff --git a/astrbot/core/knowledge_base/__init__.py b/astrbot/core/knowledge_base/__init__.py deleted file mode 100644 index df403436..00000000 --- a/astrbot/core/knowledge_base/__init__.py +++ /dev/null @@ -1,34 +0,0 @@ -""" -知识库管理模块 - -提供文档上传、解析、分块、向量化、检索等功能 -""" - -from astrbot.core.knowledge_base.models import ( - KBChunk, - KBDocument, - KBMedia, - KBSessionConfig, - KnowledgeBase, -) - -# 注意: 以下导入在对应模块实现后取消注释 -from .database import KBDatabase -from .manager import KBManager -from .manager_ops import KBManagerOps -from .session_config_db import SessionConfigDB - -# from .injector import KnowledgeBaseInjector - -__all__ = [ - "KnowledgeBase", - "KBDocument", - "KBChunk", - "KBMedia", - "KBSessionConfig", - "KBDatabase", - "SessionConfigDB", - "KBManager", - "KBManagerOps", - # "KnowledgeBaseInjector", -] diff --git a/astrbot/core/knowledge_base/chunking/base.py b/astrbot/core/knowledge_base/chunking/base.py index bcc29a5c..5aaf84ba 100644 --- a/astrbot/core/knowledge_base/chunking/base.py +++ b/astrbot/core/knowledge_base/chunking/base.py @@ -13,7 +13,7 @@ class BaseChunker(ABC): """ @abstractmethod - async def chunk(self, text: str) -> list[str]: + async def chunk(self, text: str, **kwargs) -> list[str]: """将文本分块 Args: diff --git a/astrbot/core/knowledge_base/chunking/fixed_size.py b/astrbot/core/knowledge_base/chunking/fixed_size.py index 4d1a1b28..850c5adf 100644 --- a/astrbot/core/knowledge_base/chunking/fixed_size.py +++ b/astrbot/core/knowledge_base/chunking/fixed_size.py @@ -22,7 +22,7 @@ class FixedSizeChunker(BaseChunker): self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap - async def chunk(self, text: str) -> list[str]: + async def chunk(self, text: str, **kwargs) -> list[str]: """固定大小分块 Args: @@ -31,22 +31,25 @@ class FixedSizeChunker(BaseChunker): Returns: list[str]: 分块后的文本列表 """ + chunk_size = kwargs.get("chunk_size", self.chunk_size) + chunk_overlap = kwargs.get("chunk_overlap", self.chunk_overlap) + chunks = [] start = 0 text_len = len(text) while start < text_len: - end = start + self.chunk_size + end = start + chunk_size chunk = text[start:end] if chunk: chunks.append(chunk) # 移动窗口,保留重叠部分 - start = end - self.chunk_overlap + start = end - chunk_overlap # 防止无限循环: 如果重叠过大,直接移到end - if start >= end or self.chunk_overlap >= self.chunk_size: + if start >= end or chunk_overlap >= chunk_size: start = end return chunks diff --git a/astrbot/core/knowledge_base/database.py b/astrbot/core/knowledge_base/database.py deleted file mode 100644 index 83ec0e4b..00000000 --- a/astrbot/core/knowledge_base/database.py +++ /dev/null @@ -1,347 +0,0 @@ -"""知识库数据库操作类 - -该模块封装知识库、文档、块、多媒体和会话配置相关的数据库查询操作。 - -注意: -- 该模块操作的是独立的知识库数据库 (data/knowledge_base/kb.db) -- 会话配置也存储在此数据库中，会话ID来源于主数据库 -""" - -import json -from typing import Optional - -from sqlalchemy import func, select - -from astrbot.core.knowledge_base.kb_sqlite import KBSQLiteDatabase -from astrbot.core.knowledge_base.models import ( - KBChunk, - KBDocument, - KBMedia, - KBSessionConfig, - KnowledgeBase, -) - - -class KBDatabase: - """知识库数据库操作类 - - 职责: - - 封装知识库、文档、块、多媒体和会话配置的数据库查询操作 - - 统一异常处理 - - 注意: - - 该类操作独立的知识库数据库 (kb.db) - - 会话配置存储会话ID与知识库的绑定关系，会话ID来源于主数据库 - """ - - def __init__(self, kb_db: KBSQLiteDatabase): - """初始化知识库数据库操作类 - - Args: - kb_db: 知识库独立数据库实例,而非主数据库 - """ - self.db = kb_db - - # ===== 知识库查询 ===== - - async def get_kb_by_id(self, kb_id: str) -> Optional[KnowledgeBase]: - """根据 ID 获取知识库""" - async with self.db.get_db() as session: - stmt = select(KnowledgeBase).where(KnowledgeBase.kb_id == kb_id) - result = await session.execute(stmt) - return result.scalar_one_or_none() - - async def get_kb_by_name(self, kb_name: str) -> Optional[KnowledgeBase]: - """根据名称获取知识库""" - async with self.db.get_db() as session: - stmt = select(KnowledgeBase).where(KnowledgeBase.kb_name == kb_name) - result = await session.execute(stmt) - return result.scalar_one_or_none() - - async def list_kbs(self, offset: int = 0, limit: int = 100) -> list[KnowledgeBase]: - """列出所有知识库""" - async with self.db.get_db() as session: - stmt = ( - select(KnowledgeBase) - .offset(offset) - .limit(limit) - .order_by(KnowledgeBase.created_at.desc()) - ) - result = await session.execute(stmt) - return list(result.scalars().all()) - - async def count_kbs(self) -> int: - """统计知识库数量""" - async with self.db.get_db() as session: - stmt = select(func.count(KnowledgeBase.id)) - result = await session.execute(stmt) - return result.scalar() or 0 - - # ===== 文档查询 ===== - - async def get_document_by_id(self, doc_id: str) -> Optional[KBDocument]: - """根据 ID 获取文档""" - async with self.db.get_db() as session: - stmt = select(KBDocument).where(KBDocument.doc_id == doc_id) - result = await session.execute(stmt) - return result.scalar_one_or_none() - - async def list_documents_by_kb( - self, kb_id: str, offset: int = 0, limit: int = 100 - ) -> list[KBDocument]: - """列出知识库的所有文档""" - async with self.db.get_db() as session: - stmt = ( - select(KBDocument) - .where(KBDocument.kb_id == kb_id) - .offset(offset) - .limit(limit) - .order_by(KBDocument.created_at.desc()) - ) - result = await session.execute(stmt) - return list(result.scalars().all()) - - async def count_documents_by_kb(self, kb_id: str) -> int: - """统计知识库的文档数量""" - async with self.db.get_db() as session: - stmt = select(func.count(KBDocument.id)).where(KBDocument.kb_id == kb_id) - result = await session.execute(stmt) - return result.scalar() or 0 - - # ===== 块查询 ===== - - async def get_chunk_by_id(self, chunk_id: str) -> Optional[KBChunk]: - """根据 ID 获取块""" - async with self.db.get_db() as session: - stmt = select(KBChunk).where(KBChunk.chunk_id == chunk_id) - result = await session.execute(stmt) - return result.scalar_one_or_none() - - async def get_chunks_by_kb_ids(self, kb_ids: list[str]) -> list[KBChunk]: - """根据知识库 ID 列表获取所有块""" - async with self.db.get_db() as session: - stmt = select(KBChunk).where(KBChunk.kb_id.in_(kb_ids)) - result = await session.execute(stmt) - return list(result.scalars().all()) - - async def get_chunk_by_vec_doc_id(self, vec_doc_id: str) -> Optional[KBChunk]: - """根据向量文档 ID 获取块""" - async with self.db.get_db() as session: - stmt = select(KBChunk).where(KBChunk.vec_doc_id == vec_doc_id) - result = await session.execute(stmt) - return result.scalar_one_or_none() - - async def get_chunk_with_metadata(self, chunk_id: str) -> Optional[dict]: - """获取块及其关联的文档和知识库元数据""" - async with self.db.get_db() as session: - stmt = ( - select(KBChunk, KBDocument, KnowledgeBase) - .join(KBDocument, KBChunk.doc_id == KBDocument.doc_id) - .join(KnowledgeBase, KBChunk.kb_id == KnowledgeBase.kb_id) - .where(KBChunk.chunk_id == chunk_id) - ) - result = await session.execute(stmt) - row = result.first() - - if not row: - return None - - chunk, doc, kb = row - return { - "chunk": chunk, - "document": doc, - "knowledge_base": kb, - } - - async def list_chunks_by_doc( - self, doc_id: str, offset: int = 0, limit: int = 100 - ) -> list[KBChunk]: - """列出文档的所有块""" - async with self.db.get_db() as session: - stmt = ( - select(KBChunk) - .where(KBChunk.doc_id == doc_id) - .offset(offset) - .limit(limit) - .order_by(KBChunk.chunk_index) - ) - result = await session.execute(stmt) - return list(result.scalars().all()) - - # ===== 多媒体查询 ===== - - async def list_media_by_doc(self, doc_id: str) -> list[KBMedia]: - """列出文档的所有多媒体资源""" - async with self.db.get_db() as session: - stmt = select(KBMedia).where(KBMedia.doc_id == doc_id) - result = await session.execute(stmt) - return list(result.scalars().all()) - - async def get_media_by_id(self, media_id: str) -> Optional[KBMedia]: - """根据 ID 获取多媒体资源""" - async with self.db.get_db() as session: - stmt = select(KBMedia).where(KBMedia.media_id == media_id) - result = await session.execute(stmt) - return result.scalar_one_or_none() - - # ===== 会话配置查询 ===== - - async def get_session_kb_ids(self, session_id: str) -> list[str]: - """获取会话关联的知识库 ID 列表 - - 查找顺序: - 1. 会话级别配置 (优先) - 2. 平台级别配置 - 3. 返回空列表 - - Args: - session_id: 会话ID（来自主数据库） - - Returns: - 知识库ID列表 - """ - async with self.db.get_db() as session: - # 1. 查找会话级别配置 - stmt = select(KBSessionConfig).where( - KBSessionConfig.scope == "session", - KBSessionConfig.scope_id == session_id, - ) - result = await session.execute(stmt) - config = result.scalar_one_or_none() - - if config: - return json.loads(config.kb_ids) - - # 2. 提取平台 ID (格式: platform:xxx:session_id) - parts = session_id.split(":") - if len(parts) >= 2: - platform_id = parts[0] - - # 查找平台级别配置 - stmt = select(KBSessionConfig).where( - KBSessionConfig.scope == "platform", - KBSessionConfig.scope_id == platform_id, - ) - result = await session.execute(stmt) - config = result.scalar_one_or_none() - - if config: - return json.loads(config.kb_ids) - - # 3. 无配置 - return [] - - async def set_session_kb_ids( - self, - scope: str, - scope_id: str, - kb_ids: list[str], - top_k: Optional[int] = None, - enable_rerank: Optional[bool] = None, - ) -> KBSessionConfig: - """设置会话知识库配置 - - Args: - scope: 配置范围 (session/platform) - scope_id: 范围标识 (会话 ID 或平台 ID，来自主数据库) - kb_ids: 知识库 ID 列表 - top_k: 返回结果数量 (可选) - enable_rerank: 是否启用 Rerank (可选) - - Returns: - 配置对象 - """ - async with self.db.get_db() as session: - # 查找现有配置 - stmt = select(KBSessionConfig).where( - KBSessionConfig.scope == scope, - KBSessionConfig.scope_id == scope_id, - ) - result = await session.execute(stmt) - config = result.scalar_one_or_none() - - if config: - # 更新现有配置 - config.kb_ids = json.dumps(kb_ids) - if top_k is not None: - config.top_k = top_k - if enable_rerank is not None: - config.enable_rerank = enable_rerank - else: - # 创建新配置 - config = KBSessionConfig( - scope=scope, - scope_id=scope_id, - kb_ids=json.dumps(kb_ids), - top_k=top_k, - enable_rerank=enable_rerank, - ) - session.add(config) - - await session.commit() - await session.refresh(config) - return config - - async def delete_session_kb_config(self, scope: str, scope_id: str) -> bool: - """删除会话知识库配置 - - Args: - scope: 配置范围 (session/platform) - scope_id: 范围标识 (会话 ID 或平台 ID) - - Returns: - 是否删除成功 - """ - async with self.db.get_db() as session: - stmt = select(KBSessionConfig).where( - KBSessionConfig.scope == scope, - KBSessionConfig.scope_id == scope_id, - ) - result = await session.execute(stmt) - config = result.scalar_one_or_none() - - if not config: - return False - - await session.delete(config) - await session.commit() - return True - - async def delete_session_kb_config_by_session_id(self, session_id: str) -> bool: - """根据会话ID删除会话配置（用于主数据库会话删除时的级联清理） - - Args: - session_id: 会话ID（来自主数据库） - - Returns: - 是否删除成功 - """ - return await self.delete_session_kb_config("session", session_id) - - async def list_all_session_configs( - self, offset: int = 0, limit: int = 100, scope: Optional[str] = None - ) -> list[KBSessionConfig]: - """列出所有会话配置 - - Args: - offset: 偏移量 - limit: 限制数量 - scope: 可选的范围过滤 (session/platform) - - Returns: - 会话配置列表 - """ - async with self.db.get_db() as session: - stmt = select(KBSessionConfig) - - if scope: - stmt = stmt.where(KBSessionConfig.scope == scope) - - stmt = ( - stmt.offset(offset) - .limit(limit) - .order_by(KBSessionConfig.created_at.desc()) - ) - - result = await session.execute(stmt) - return list(result.scalars().all()) diff --git a/astrbot/core/knowledge_base/injector.py b/astrbot/core/knowledge_base/injector.py deleted file mode 100644 index 968e51ff..00000000 --- a/astrbot/core/knowledge_base/injector.py +++ /dev/null @@ -1,115 +0,0 @@ -"""知识库上下文注入器 - -负责检索相关知识并格式化为 LLM 可用的上下文文本 -""" - -from typing import List, Optional - -from astrbot.core.knowledge_base.database import KBDatabase -from astrbot.core.knowledge_base.retrieval.manager import ( - RetrievalManager, - RetrievalResult, -) - - -class KnowledgeBaseInjector: - """知识库上下文注入器 - - 职责: - - 检索相关知识 - - 格式化为上下文文本 - - 注入到 LLM Prompt - """ - - def __init__( - self, - kb_db: KBDatabase, - retrieval_manager: RetrievalManager, - ): - """初始化知识库上下文注入器 - - Args: - kb_db: 知识库数据库实例 - retrieval_manager: 检索管理器实例 - """ - self.kb_db = kb_db - self.retrieval_manager = retrieval_manager - - async def retrieve_and_inject( - self, - unified_msg_origin: str, - query: str, - top_k: int = 5, - ) -> Optional[dict]: - """检索并注入知识库上下文 - - Args: - unified_msg_origin: 统一消息来源 ID (会话 ID) - query: 用户查询 - top_k: 返回结果数量 - - Returns: - Optional[dict]: 包含检索结果和格式化上下文的字典,如果无结果则返回 None - { - "context_text": str, # 格式化的上下文文本 - "results": List[dict], # 原始检索结果列表 - } - """ - # 1. 获取会话关联的知识库 - kb_ids = await self.kb_db.get_session_kb_ids(unified_msg_origin) - - if not kb_ids: - return None - - # 2. 检索知识 - results = await self.retrieval_manager.retrieve( - query=query, - kb_ids=kb_ids, - top_m_final=top_k, - ) - - if not results: - return None - - # 3. 格式化上下文 - context_text = self._format_context(results) - - # 4. 转换结果为字典格式 - results_dict = [ - { - "chunk_id": r.chunk_id, - "doc_id": r.doc_id, - "kb_id": r.kb_id, - "kb_name": r.kb_name, - "doc_name": r.doc_name, - "chunk_index": r.metadata.get("chunk_index", 0), - "content": r.content, - "score": r.score, - } - for r in results - ] - - return { - "context_text": context_text, - "results": results_dict, - } - - def _format_context(self, results: List[RetrievalResult]) -> str: - """格式化知识上下文 - - Args: - results: 检索结果列表 - - Returns: - str: 格式化的上下文文本 - """ - lines = ["以下是相关的知识库内容,请参考这些信息回答用户的问题:\n"] - - for i, result in enumerate(results, 1): - lines.append(f"【知识 {i}】") - lines.append(f"来源: {result.kb_name} / {result.doc_name}") - lines.append(f"内容: {result.content}") - lines.append(f"相关度: {result.score:.2f}") - lines.append("") - - return "\n".join(lines) diff --git a/astrbot/core/knowledge_base/kb_db_sqlite.py b/astrbot/core/knowledge_base/kb_db_sqlite.py new file mode 100644 index 00000000..827d621d --- /dev/null +++ b/astrbot/core/knowledge_base/kb_db_sqlite.py @@ -0,0 +1,299 @@ +from contextlib import asynccontextmanager +from pathlib import Path + +from sqlmodel import col, desc +from sqlalchemy import text, func, select, update, delete +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from astrbot.core import logger +from astrbot.core.knowledge_base.models import ( + BaseKBModel, + KBDocument, + KBMedia, + KnowledgeBase, +) +from astrbot.core.db.vec_db.faiss_impl import FaissVecDB + + +class KBSQLiteDatabase: + def __init__(self, db_path: str = "data/knowledge_base/kb.db") -> None: + """初始化知识库数据库 + + Args: + db_path: 数据库文件路径, 默认为 data/knowledge_base/kb.db + """ + self.db_path = db_path + self.DATABASE_URL = f"sqlite+aiosqlite:///{db_path}" + self.inited = False + + # 确保目录存在 + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + + # 创建异步引擎 + self.engine = create_async_engine( + self.DATABASE_URL, + echo=False, + pool_pre_ping=True, + pool_recycle=3600, + ) + + # 创建会话工厂 + self.async_session = async_sessionmaker( + self.engine, + class_=AsyncSession, + expire_on_commit=False, + ) + + @asynccontextmanager + async def get_db(self): + """获取数据库会话 + + 用法: + async with kb_db.get_db() as session: + # 执行数据库操作 + result = await session.execute(stmt) + """ + async with self.async_session() as session: + yield session + + async def initialize(self) -> None: + """初始化数据库,创建表并配置 SQLite 参数""" + async with self.engine.begin() as conn: + # 创建所有知识库相关表 + await conn.run_sync(BaseKBModel.metadata.create_all) + + # 配置 SQLite 性能优化参数 + await conn.execute(text("PRAGMA journal_mode=WAL")) + await conn.execute(text("PRAGMA synchronous=NORMAL")) + await conn.execute(text("PRAGMA cache_size=20000")) + await conn.execute(text("PRAGMA temp_store=MEMORY")) + await conn.execute(text("PRAGMA mmap_size=134217728")) + await conn.execute(text("PRAGMA optimize")) + await conn.commit() + + self.inited = True + + async def migrate_to_v1(self) -> None: + """执行知识库数据库 v1 迁移 + + 创建所有必要的索引以优化查询性能 + """ + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + # 创建知识库表索引 + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_kb_kb_id " + "ON knowledge_bases(kb_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_kb_name " + "ON knowledge_bases(kb_name)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_kb_created_at " + "ON knowledge_bases(created_at)" + ) + ) + + # 创建文档表索引 + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_doc_id " + "ON kb_documents(doc_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_kb_id " + "ON kb_documents(kb_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_name " + "ON kb_documents(doc_name)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_type " + "ON kb_documents(file_type)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_created_at " + "ON kb_documents(created_at)" + ) + ) + + # 创建多媒体表索引 + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_media_id " + "ON kb_media(media_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_doc_id " + "ON kb_media(doc_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_kb_id ON kb_media(kb_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_type " + "ON kb_media(media_type)" + ) + ) + + await session.commit() + + async def close(self) -> None: + """关闭数据库连接""" + await self.engine.dispose() + logger.info(f"知识库数据库已关闭: {self.db_path}") + + async def get_kb_by_id(self, kb_id: str) -> KnowledgeBase | None: + """根据 ID 获取知识库""" + async with self.get_db() as session: + stmt = select(KnowledgeBase).where(col(KnowledgeBase.kb_id) == kb_id) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def get_kb_by_name(self, kb_name: str) -> KnowledgeBase | None: + """根据名称获取知识库""" + async with self.get_db() as session: + stmt = select(KnowledgeBase).where(col(KnowledgeBase.kb_name) == kb_name) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def list_kbs(self, offset: int = 0, limit: int = 100) -> list[KnowledgeBase]: + """列出所有知识库""" + async with self.get_db() as session: + stmt = ( + select(KnowledgeBase) + .offset(offset) + .limit(limit) + .order_by(desc(KnowledgeBase.created_at)) + ) + result = await session.execute(stmt) + return list(result.scalars().all()) + + async def count_kbs(self) -> int: + """统计知识库数量""" + async with self.get_db() as session: + stmt = select(func.count(col(KnowledgeBase.id))) + result = await session.execute(stmt) + return result.scalar() or 0 + + # ===== 文档查询 ===== + + async def get_document_by_id(self, doc_id: str) -> KBDocument | None: + """根据 ID 获取文档""" + async with self.get_db() as session: + stmt = select(KBDocument).where(col(KBDocument.doc_id) == doc_id) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def list_documents_by_kb( + self, kb_id: str, offset: int = 0, limit: int = 100 + ) -> list[KBDocument]: + """列出知识库的所有文档""" + async with self.get_db() as session: + stmt = ( + select(KBDocument) + .where(col(KBDocument.kb_id) == kb_id) + .offset(offset) + .limit(limit) + .order_by(desc(KBDocument.created_at)) + ) + result = await session.execute(stmt) + return list(result.scalars().all()) + + async def count_documents_by_kb(self, kb_id: str) -> int: + """统计知识库的文档数量""" + async with self.get_db() as session: + stmt = select(func.count(col(KBDocument.id))).where( + col(KBDocument.kb_id) == kb_id + ) + result = await session.execute(stmt) + return result.scalar() or 0 + + async def get_document_with_metadata(self, doc_id: str) -> dict | None: + async with self.get_db() as session: + stmt = ( + select(KBDocument, KnowledgeBase) + .join(KnowledgeBase, col(KBDocument.kb_id) == col(KnowledgeBase.kb_id)) + .where(col(KBDocument.doc_id) == doc_id) + ) + result = await session.execute(stmt) + row = result.first() + + if not row: + return None + + return { + "document": row[0], + "knowledge_base": row[1], + } + + async def delete_document_by_id(self, doc_id: str, vec_db: FaissVecDB): + """删除单个文档及其相关数据""" + # 在知识库表中删除 + async with self.get_db() as session: + async with session.begin(): + # 删除文档记录 + delete_stmt = delete(KBDocument).where(col(KBDocument.doc_id) == doc_id) + await session.execute(delete_stmt) + await session.commit() + + # 在 vec db 中删除相关向量 + await vec_db.delete_documents(metadata_filters={"kb_doc_id": doc_id}) + + # ===== 多媒体查询 ===== + + async def list_media_by_doc(self, doc_id: str) -> list[KBMedia]: + """列出文档的所有多媒体资源""" + async with self.get_db() as session: + stmt = select(KBMedia).where(col(KBMedia.doc_id) == doc_id) + result = await session.execute(stmt) + return list(result.scalars().all()) + + async def get_media_by_id(self, media_id: str) -> KBMedia | None: + """根据 ID 获取多媒体资源""" + async with self.get_db() as session: + stmt = select(KBMedia).where(col(KBMedia.media_id) == media_id) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def update_kb_stats(self, kb_id: str, vec_db: FaissVecDB) -> None: + """更新知识库统计信息""" + chunk_cnt = await vec_db.count_documents() + + async with self.get_db() as session: + async with session.begin(): + update_stmt = ( + update(KnowledgeBase) + .where(col(KnowledgeBase.kb_id) == kb_id) + .values( + doc_count=select(func.count(col(KBDocument.id))) + .where(col(KBDocument.kb_id) == kb_id) + .scalar_subquery(), + chunk_count=chunk_cnt, + ) + ) + + await session.execute(update_stmt) + await session.commit() diff --git a/astrbot/core/knowledge_base/kb_helper.py b/astrbot/core/knowledge_base/kb_helper.py new file mode 100644 index 00000000..dc9febea --- /dev/null +++ b/astrbot/core/knowledge_base/kb_helper.py @@ -0,0 +1,352 @@ +import uuid +import aiofiles +import json +from pathlib import Path +from .models import KnowledgeBase, KBDocument, KBMedia +from .kb_db_sqlite import KBSQLiteDatabase +from astrbot.core.db.vec_db.base import BaseVecDB +from astrbot.core.db.vec_db.faiss_impl.vec_db import FaissVecDB +from astrbot.core.provider.provider import EmbeddingProvider, RerankProvider +from astrbot.core.provider.manager import ProviderManager +from .parsers.base import BaseParser +from .chunking.base import BaseChunker +from astrbot.core import logger + + +class KBHelper: + vec_db: BaseVecDB + kb: KnowledgeBase + + def __init__( + self, + kb_db: KBSQLiteDatabase, + kb: KnowledgeBase, + provider_manager: ProviderManager, + kb_root_dir: str, + chunker: BaseChunker, + parsers: dict[str, BaseParser], + ): + self.kb_db = kb_db + self.kb = kb + self.prov_mgr = provider_manager + self.kb_root_dir = kb_root_dir + self.parsers = parsers + self.chunker = chunker + + self.kb_dir = Path(self.kb_root_dir) / self.kb.kb_id + self.kb_medias_dir = Path(self.kb_dir) / "medias" / self.kb.kb_id + self.kb_files_dir = Path(self.kb_dir) / "files" / self.kb.kb_id + + self.kb_medias_dir.mkdir(parents=True, exist_ok=True) + self.kb_files_dir.mkdir(parents=True, exist_ok=True) + + async def initialize(self): + await self._ensure_vec_db() + + async def get_ep(self) -> EmbeddingProvider: + if not self.kb.embedding_provider_id: + raise ValueError(f"知识库 {self.kb.kb_name} 未配置 Embedding Provider") + ep: EmbeddingProvider = await self.prov_mgr.get_provider_by_id( + self.kb.embedding_provider_id + ) # type: ignore + if not ep: + raise ValueError( + f"无法找到 ID 为 {self.kb.embedding_provider_id} 的 Embedding Provider" + ) + return ep + + async def get_rp(self) -> RerankProvider | None: + if not self.kb.rerank_provider_id: + return None + rp: RerankProvider = await self.prov_mgr.get_provider_by_id( + self.kb.rerank_provider_id + ) # type: ignore + if not rp: + raise ValueError( + f"无法找到 ID 为 {self.kb.rerank_provider_id} 的 Rerank Provider" + ) + return rp + + async def _ensure_vec_db(self) -> FaissVecDB: + if not self.kb.embedding_provider_id: + raise ValueError(f"知识库 {self.kb.kb_name} 未配置 Embedding Provider") + + ep = await self.get_ep() + rp = await self.get_rp() + + vec_db = FaissVecDB( + doc_store_path=str(self.kb_dir / "doc.db"), + index_store_path=str(self.kb_dir / "index.faiss"), + embedding_provider=ep, + rerank_provider=rp, + ) + await vec_db.initialize() + self.vec_db = vec_db + return vec_db + + async def delete_vec_db(self): + await self.terminate() + if self.kb_dir.exists(): + for item in self.kb_dir.iterdir(): + if item.is_file(): + item.unlink() + self.kb_dir.rmdir() + + async def terminate(self): + if self.vec_db: + await self.vec_db.close() + + async def upload_document( + self, + file_name: str, + file_content: bytes, + file_type: str, + chunk_size: int = 512, + chunk_overlap: int = 50, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> KBDocument: + """上传并处理文档（带原子性保证和失败清理） + + 流程: + 1. 保存原始文件 + 2. 解析文档内容 + 3. 提取多媒体资源 + 4. 分块处理 + 5. 生成向量并存储 + 6. 保存元数据（事务） + 7. 更新统计 + + Args: + progress_callback: 进度回调函数，接收参数 (stage, current, total) + - stage: 当前阶段 ('parsing', 'chunking', 'embedding') + - current: 当前进度 + - total: 总数 + """ + await self._ensure_vec_db() + doc_id = str(uuid.uuid4()) + media_paths: list[Path] = [] + + # file_path = self.kb_files_dir / f"{doc_id}.{file_type}" + # async with aiofiles.open(file_path, "wb") as f: + # await f.write(file_content) + + try: + # 阶段1: 解析文档 + if progress_callback: + await progress_callback("parsing", 0, 100) + + parser = self.parsers.get(file_type) + if not parser: + raise ValueError(f"不支持的文件类型: {file_type}") + parse_result = await parser.parse(file_content, file_name) + text_content = parse_result.text + media_items = parse_result.media + + if progress_callback: + await progress_callback("parsing", 100, 100) + + # 保存媒体文件 + saved_media = [] + for media_item in media_items: + media = await self._save_media( + doc_id=doc_id, + media_type=media_item.media_type, + file_name=media_item.file_name, + content=media_item.content, + mime_type=media_item.mime_type, + ) + saved_media.append(media) + media_paths.append(Path(media.file_path)) + + # 阶段2: 分块 + if progress_callback: + await progress_callback("chunking", 0, 100) + + chunks_text = await self.chunker.chunk( + text_content, chunk_size=chunk_size, chunk_overlap=chunk_overlap + ) + contents = [] + metadatas = [] + for idx, chunk_text in enumerate(chunks_text): + contents.append(chunk_text) + metadatas.append( + { + "kb_id": self.kb.kb_id, + "kb_doc_id": doc_id, + "chunk_index": idx, + } + ) + + if progress_callback: + await progress_callback("chunking", 100, 100) + + # 阶段3: 生成向量（带进度回调） + async def embedding_progress_callback(current, total): + if progress_callback: + await progress_callback("embedding", current, total) + + await self.vec_db.insert_batch( + contents=contents, + metadatas=metadatas, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=embedding_progress_callback, + ) + + # 保存文档的元数据 + doc = KBDocument( + doc_id=doc_id, + kb_id=self.kb.kb_id, + doc_name=file_name, + file_type=file_type, + file_size=len(file_content), + # file_path=str(file_path), + file_path="", + chunk_count=len(chunks_text), + media_count=0, + ) + async with self.kb_db.get_db() as session: + async with session.begin(): + session.add(doc) + for media in saved_media: + session.add(media) + await session.commit() + + await session.refresh(doc) + + vec_db: FaissVecDB = self.vec_db # type: ignore + await self.kb_db.update_kb_stats(kb_id=self.kb.kb_id, vec_db=vec_db) + await self.refresh_kb() + await self.refresh_document(doc_id) + return doc + except Exception as e: + logger.error(f"上传文档失败: {e}") + # if file_path.exists(): + # file_path.unlink() + + for media_path in media_paths: + try: + if media_path.exists(): + media_path.unlink() + except Exception as me: + logger.warning(f"清理多媒体文件失败 {media_path}: {me}") + + raise e + + async def list_documents( + self, offset: int = 0, limit: int = 100 + ) -> list[KBDocument]: + """列出知识库的所有文档""" + docs = await self.kb_db.list_documents_by_kb(self.kb.kb_id, offset, limit) + return docs + + async def get_document(self, doc_id: str) -> KBDocument | None: + """获取单个文档""" + doc = await self.kb_db.get_document_by_id(doc_id) + return doc + + async def delete_document(self, doc_id: str): + """删除单个文档及其相关数据""" + await self.kb_db.delete_document_by_id( + doc_id=doc_id, + vec_db=self.vec_db, # type: ignore + ) + await self.kb_db.update_kb_stats( + kb_id=self.kb.kb_id, + vec_db=self.vec_db, # type: ignore + ) + await self.refresh_kb() + + async def delete_chunk(self, chunk_id: str, doc_id: str): + """删除单个文本块及其相关数据""" + vec_db: FaissVecDB = self.vec_db # type: ignore + await vec_db.delete(chunk_id) + await self.kb_db.update_kb_stats( + kb_id=self.kb.kb_id, + vec_db=self.vec_db, # type: ignore + ) + await self.refresh_kb() + await self.refresh_document(doc_id) + + async def refresh_kb(self): + if self.kb: + kb = await self.kb_db.get_kb_by_id(self.kb.kb_id) + if kb: + self.kb = kb + + async def refresh_document(self, doc_id: str) -> None: + """更新文档的元数据""" + doc = await self.get_document(doc_id) + if not doc: + raise ValueError(f"无法找到 ID 为 {doc_id} 的文档") + chunk_count = await self.get_chunk_count_by_doc_id(doc_id) + doc.chunk_count = chunk_count + async with self.kb_db.get_db() as session: + async with session.begin(): + session.add(doc) + await session.commit() + await session.refresh(doc) + + async def get_chunks_by_doc_id( + self, doc_id: str, offset: int = 0, limit: int = 100 + ) -> list[dict]: + """获取文档的所有块及其元数据""" + vec_db: FaissVecDB = self.vec_db # type: ignore + chunks = await vec_db.document_storage.get_documents( + metadata_filters={"kb_doc_id": doc_id}, offset=offset, limit=limit + ) + result = [] + for chunk in chunks: + chunk_md = json.loads(chunk["metadata"]) + result.append( + { + "chunk_id": chunk["doc_id"], + "doc_id": chunk_md["kb_doc_id"], + "kb_id": chunk_md["kb_id"], + "chunk_index": chunk_md["chunk_index"], + "content": chunk["text"], + "char_count": len(chunk["text"]), + } + ) + return result + + async def get_chunk_count_by_doc_id(self, doc_id: str) -> int: + """获取文档的块数量""" + vec_db: FaissVecDB = self.vec_db # type: ignore + count = await vec_db.count_documents(metadata_filter={"kb_doc_id": doc_id}) + return count + + async def _save_media( + self, + doc_id: str, + media_type: str, + file_name: str, + content: bytes, + mime_type: str, + ) -> KBMedia: + """保存多媒体资源""" + media_id = str(uuid.uuid4()) + ext = Path(file_name).suffix + + # 保存文件 + file_path = self.kb_medias_dir / doc_id / f"{media_id}{ext}" + file_path.parent.mkdir(parents=True, exist_ok=True) + async with aiofiles.open(file_path, "wb") as f: + await f.write(content) + + media = KBMedia( + media_id=media_id, + doc_id=doc_id, + kb_id=self.kb.kb_id, + media_type=media_type, + file_name=file_name, + file_path=str(file_path), + file_size=len(content), + mime_type=mime_type, + ) + + return media diff --git a/astrbot/core/knowledge_base/kb_manager_lifecycle.py b/astrbot/core/knowledge_base/kb_manager_lifecycle.py deleted file mode 100644 index 0768874f..00000000 --- a/astrbot/core/knowledge_base/kb_manager_lifecycle.py +++ /dev/null @@ -1,364 +0,0 @@ -""" -知识库管理器 -负责知识库模块的初始化、配置和资源管理 - -架构说明: -- 知识库数据存储在独立的数据库 (data/knowledge_base/kb.db) -- 会话配置存储在主数据库 (data/astrbot.db) 以便于会话关联 -""" - -from pathlib import Path -from astrbot.core import logger -from astrbot.core.db import BaseDatabase -from astrbot.core.provider.manager import ProviderManager - - -class KnowledgeBaseManager: - """知识库管理器 - - 职责: - - 知识库模块的初始化 - - Embedding Provider 和 Rerank Provider 的选择 - - 各个子组件的协调管理 - - 注册会话删除回调，实现级联清理 - - 架构说明: - - 知识库数据存储在独立数据库 (kb.db) - - 会话配置存储在独立数据库 (kb.db)，会话ID来自主数据库 - - 通过回调机制实现与主数据库的生命周期同步 - """ - - def __init__( - self, - config: dict, - main_db: BaseDatabase, - provider_manager: ProviderManager, - ): - """初始化知识库管理器 - - Args: - config: 配置字典 - main_db: 主数据库实例 (不直接使用，仅用于类型引用) - provider_manager: Provider 管理器 - """ - self.config = config.get("knowledge_base", {}) - self.provider_manager = provider_manager - - # 知识库独立数据库 - self.kb_db = None - - # 组件实例 - self.kb_database = None - self.kb_manager = None - self.kb_vec_db = None - self.retrieval_manager = None - self.kb_injector = None - - self._initialized = False - self._session_deleted_callback_registered = False - - async def initialize(self): - """初始化知识库模块""" - if not self.config.get("enabled", False): - logger.info("知识库功能未启用") - return - - try: - logger.info("正在初始化知识库模块...") - - # 1. 检查并选择 Embedding Provider - embedding_provider = self._select_embedding_provider() - if not embedding_provider: - logger.warning("未配置 Embedding Provider，知识库功能无法使用") - return - - # 2. 初始化数据库 - await self._init_kb_database() - await self._init_database() - - # 3. 初始化向量数据库 - await self._init_vector_db(embedding_provider) - - # 4. 初始化解析器和分块器 - parsers = self._init_parsers() - chunker = self._init_chunker() - - # 5. 初始化知识库管理器 - await self._init_kb_manager(parsers, chunker) - - # 6. 初始化检索管理器 - await self._init_retrieval_manager() - - # 7. 初始化上下文注入器 - await self._init_injector() - - self._initialized = True - logger.info("知识库模块初始化完成") - - except ImportError as e: - logger.error(f"知识库模块导入失败: {e}") - logger.warning("请确保已安装所需依赖: pypdf, aiofiles, Pillow, rank-bm25") - except Exception as e: - logger.error(f"知识库模块初始化失败: {e}") - import traceback - - logger.error(traceback.format_exc()) - - async def _init_kb_database(self): - """初始化知识库独立数据库""" - from astrbot.core.knowledge_base.kb_sqlite import KBSQLiteDatabase - - db_path = self.config.get("storage", {}).get( - "kb_db_path", "data/knowledge_base/kb.db" - ) - Path(db_path).parent.mkdir(parents=True, exist_ok=True) - - self.kb_db = KBSQLiteDatabase(db_path) - await self.kb_db.initialize() - await self.kb_db.migrate_to_v1() - - logger.info(f"知识库独立数据库已初始化: {db_path}") - - async def _init_database(self): - """初始化知识库数据库操作类""" - from astrbot.core.knowledge_base.database import KBDatabase - - self.kb_database = KBDatabase(self.kb_db) - - async def _init_vector_db(self, embedding_provider): - """初始化向量数据库""" - from astrbot.core.db.vec_db.faiss_impl import FaissVecDB - - storage_path = self.config.get("storage", {}).get( - "vector_db_path", "data/knowledge_base/vectors" - ) - Path(storage_path).mkdir(parents=True, exist_ok=True) - - self.kb_vec_db = FaissVecDB( - doc_store_path=f"{storage_path}/documents.db", - index_store_path=f"{storage_path}/index.faiss", - embedding_provider=embedding_provider, - ) - await self.kb_vec_db.initialize() - - def _init_parsers(self) -> dict: - """初始化文档解析器""" - from astrbot.core.knowledge_base.parsers.text_parser import TextParser - from astrbot.core.knowledge_base.parsers.pdf_parser import PDFParser - - return { - "txt": TextParser(), - "md": TextParser(), - "markdown": TextParser(), - "pdf": PDFParser(), - } - - def _init_chunker(self): - """初始化分块器""" - from astrbot.core.knowledge_base.chunking.fixed_size import FixedSizeChunker - - chunking_config = self.config.get("chunking", {}) - return FixedSizeChunker( - chunk_size=chunking_config.get("chunk_size", 512), - chunk_overlap=chunking_config.get("chunk_overlap", 50), - ) - - async def _init_kb_manager(self, parsers: dict, chunker): - """初始化知识库管理器""" - from astrbot.core.knowledge_base.manager import KBManager - - files_path = self.config.get("storage", {}).get( - "files_path", "data/knowledge_base" - ) - - self.kb_manager = KBManager( - db=self.kb_db, # 使用独立的知识库数据库 - vec_db=self.kb_vec_db, - storage_path=files_path, - parsers=parsers, - chunker=chunker, - provider_manager=self.provider_manager, - ) - - async def _init_retrieval_manager(self): - """初始化检索管理器""" - from astrbot.core.knowledge_base.retrieval.manager import RetrievalManager - from astrbot.core.knowledge_base.retrieval.sparse_retriever import ( - SparseRetriever, - ) - from astrbot.core.knowledge_base.retrieval.rank_fusion import RankFusion - - sparse_retriever = SparseRetriever(self.kb_database) - rank_fusion = RankFusion(self.kb_database) - - # 选择 Rerank Provider (可选) - rerank_provider = self._select_rerank_provider() - - self.retrieval_manager = RetrievalManager( - vec_db=self.kb_vec_db, - sparse_retriever=sparse_retriever, - rank_fusion=rank_fusion, - kb_db=self.kb_database, - rerank_provider=rerank_provider, - ) - - async def _init_injector(self): - """初始化上下文注入器""" - from astrbot.core.knowledge_base.injector import KnowledgeBaseInjector - - self.kb_injector = KnowledgeBaseInjector( - kb_db=self.kb_database, - retrieval_manager=self.retrieval_manager, - ) - - def _select_embedding_provider(self): - """选择 Embedding Provider - - 逻辑: - - 如果配置了 embedding_provider_id，则使用指定的 provider - - 如果没有配置，但有 embedding provider，则使用第一个 - - 如果有多个 embedding provider 但没有指定，则警告并使用第一个 - """ - embedding_providers = self.provider_manager.embedding_provider_insts - - if not embedding_providers: - return None - - configured_provider_id = self.config.get("embedding_provider_id") - - if configured_provider_id: - # 按 ID 查找 - for provider in embedding_providers: - provider_id = provider.meta().id - if provider_id == configured_provider_id: - logger.info(f"知识库使用 Embedding Provider: {provider_id}") - return provider - logger.warning( - f"未找到配置的 Embedding Provider ID: {configured_provider_id}，" - f"将使用第一个可用的" - ) - - if len(embedding_providers) > 1 and not configured_provider_id: - provider = embedding_providers[0] - provider_id = provider.meta().id - logger.info( - f"检测到 {len(embedding_providers)} 个 Embedding Provider，" - f"未在配置文件中指定 embedding_provider_id，将使用第一个: {provider_id}" - ) - return provider - - provider = embedding_providers[0] - provider_id = provider.meta().id - logger.info(f"知识库使用 Embedding Provider: {provider_id}") - return provider - - def _select_rerank_provider(self): - """选择 Rerank Provider (可选)""" - if not self.config.get("retrieval", {}).get("enable_rerank", True): - return None - - rerank_providers = self.provider_manager.rerank_provider_insts - if not rerank_providers: - return None - - configured_provider_id = self.config.get("rerank_provider_id") - - if configured_provider_id: - for provider in rerank_providers: - provider_id = provider.meta().id - if provider_id == configured_provider_id: - logger.info(f"知识库使用 Rerank Provider: {provider_id}") - return provider - logger.warning(f"未找到配置的 Rerank Provider ID: {configured_provider_id}") - - if len(rerank_providers) > 0: - provider = rerank_providers[0] - provider_id = provider.meta().id - logger.info(f"知识库使用 Rerank Provider: {provider_id}") - return provider - - return None - - @property - def is_initialized(self) -> bool: - """检查是否已初始化""" - return self._initialized - - def get_kb_manager(self): - """获取知识库管理器""" - return self.kb_manager if self._initialized else None - - def get_kb_injector(self): - """获取知识库上下文注入器""" - return self.kb_injector if self._initialized else None - - def register_session_lifecycle_hooks(self, conversation_manager): - """注册会话生命周期钩子 - - 在会话删除时自动清理知识库配置，实现零侵入的级联清理。 - - Args: - conversation_manager: 会话管理器实例 - """ - if self._session_deleted_callback_registered or not self._initialized: - return - - async def on_session_deleted(session_id: str): - """会话删除回调：清理知识库配置""" - try: - await self.kb_database.delete_session_kb_config_by_session_id( - session_id - ) - logger.info(f"已清理会话知识库配置: {session_id}") - except Exception as e: - logger.error(f"清理会话知识库配置失败 ({session_id}): {e}") - - conversation_manager.register_on_session_deleted(on_session_deleted) - self._session_deleted_callback_registered = True - logger.info("已注册知识库会话删除回调") - - async def reinitialize(self): - """重新初始化知识库模块 - - 用于在运行时动态初始化知识库模块（例如用户添加了 embedding provider 后） - """ - if self._initialized: - logger.info("知识库模块已初始化，将重新初始化") - await self.terminate() - - await self.initialize() - return self._initialized - - async def terminate(self): - """终止知识库模块，清理资源""" - if not self._initialized: - return - - logger.info("正在终止知识库模块...") - - # 关闭向量数据库连接 - if self.kb_vec_db: - try: - await self.kb_vec_db.close() - logger.debug("向量数据库已关闭") - except Exception as e: - logger.warning(f"关闭向量数据库时出错: {e}") - - # 关闭知识库独立数据库连接 - if self.kb_db: - try: - await self.kb_db.close() - logger.debug("知识库数据库已关闭") - except Exception as e: - logger.warning(f"关闭知识库数据库时出错: {e}") - - # 清理资源 - self._initialized = False - self.kb_db = None - self.kb_database = None - self.kb_manager = None - self.kb_vec_db = None - self.retrieval_manager = None - self.kb_injector = None - - logger.info("知识库模块已终止") diff --git a/astrbot/core/knowledge_base/kb_mgr.py b/astrbot/core/knowledge_base/kb_mgr.py new file mode 100644 index 00000000..0a9c16ad --- /dev/null +++ b/astrbot/core/knowledge_base/kb_mgr.py @@ -0,0 +1,279 @@ +import traceback +from pathlib import Path +from astrbot.core import logger +from astrbot.core.provider.manager import ProviderManager + +from .retrieval.manager import RetrievalManager, RetrievalResult +from .retrieval.sparse_retriever import SparseRetriever +from .retrieval.rank_fusion import RankFusion +from .kb_db_sqlite import KBSQLiteDatabase + +from .parsers.text_parser import TextParser +from .parsers.pdf_parser import PDFParser +from .chunking.fixed_size import FixedSizeChunker +from .kb_helper import KBHelper + +from .models import KnowledgeBase + + +FILES_PATH = "data/knowledge_base" +DB_PATH = Path(FILES_PATH) / "kb.db" +"""Knowledge Base storage root directory""" +PARSERS = { + "txt": TextParser(), + "md": TextParser(), + "markdown": TextParser(), + "pdf": PDFParser(), +} +CHUNKER = FixedSizeChunker() + + +class KnowledgeBaseManager: + kb_db: KBSQLiteDatabase + retrieval_manager: RetrievalManager + + def __init__( + self, + provider_manager: ProviderManager, + ): + Path(DB_PATH).parent.mkdir(parents=True, exist_ok=True) + self.provider_manager = provider_manager + self._session_deleted_callback_registered = False + + self.kb_insts: dict[str, KBHelper] = {} + + async def initialize(self): + """初始化知识库模块""" + try: + logger.info("正在初始化知识库模块...") + + # 初始化数据库 + await self._init_kb_database() + + # 初始化检索管理器 + sparse_retriever = SparseRetriever(self.kb_db) + rank_fusion = RankFusion(self.kb_db) + self.retrieval_manager = RetrievalManager( + sparse_retriever=sparse_retriever, + rank_fusion=rank_fusion, + kb_db=self.kb_db, + ) + await self.load_kbs() + + except ImportError as e: + logger.error(f"知识库模块导入失败: {e}") + logger.warning("请确保已安装所需依赖: pypdf, aiofiles, Pillow, rank-bm25") + except Exception as e: + logger.error(f"知识库模块初始化失败: {e}") + logger.error(traceback.format_exc()) + + async def _init_kb_database(self): + self.kb_db = KBSQLiteDatabase(DB_PATH.as_posix()) + await self.kb_db.initialize() + await self.kb_db.migrate_to_v1() + logger.info(f"KnowledgeBase database initialized: {DB_PATH}") + + async def load_kbs(self): + """加载所有知识库实例""" + kb_records = await self.kb_db.list_kbs() + for record in kb_records: + kb_helper = KBHelper( + kb_db=self.kb_db, + kb=record, + provider_manager=self.provider_manager, + kb_root_dir=FILES_PATH, + chunker=CHUNKER, + parsers=PARSERS, + ) + await kb_helper.initialize() + self.kb_insts[record.kb_id] = kb_helper + + async def create_kb( + self, + kb_name: str, + description: str | None = None, + emoji: str | None = None, + embedding_provider_id: str | None = None, + rerank_provider_id: str | None = None, + chunk_size: int | None = None, + chunk_overlap: int | None = None, + top_k_dense: int | None = None, + top_k_sparse: int | None = None, + top_m_final: int | None = None, + ) -> KBHelper: + """创建新的知识库实例""" + kb = KnowledgeBase( + kb_name=kb_name, + description=description, + emoji=emoji or "📚", + embedding_provider_id=embedding_provider_id, + rerank_provider_id=rerank_provider_id, + chunk_size=chunk_size if chunk_size is not None else 512, + chunk_overlap=chunk_overlap if chunk_overlap is not None else 50, + top_k_dense=top_k_dense if top_k_dense is not None else 50, + top_k_sparse=top_k_sparse if top_k_sparse is not None else 50, + top_m_final=top_m_final if top_m_final is not None else 5, + ) + async with self.kb_db.get_db() as session: + session.add(kb) + await session.commit() + await session.refresh(kb) + + kb_helper = KBHelper( + kb_db=self.kb_db, + kb=kb, + provider_manager=self.provider_manager, + kb_root_dir=FILES_PATH, + chunker=CHUNKER, + parsers=PARSERS, + ) + await kb_helper.initialize() + self.kb_insts[kb.kb_id] = kb_helper + return kb_helper + + async def get_kb(self, kb_id: str) -> KBHelper | None: + """获取知识库实例""" + if kb_id in self.kb_insts: + return self.kb_insts[kb_id] + + async def get_kb_by_name(self, kb_name: str) -> KBHelper | None: + """通过名称获取知识库实例""" + for kb_helper in self.kb_insts.values(): + if kb_helper.kb.kb_name == kb_name: + return kb_helper + return None + + async def delete_kb(self, kb_id: str) -> bool: + """删除知识库实例""" + kb_helper = await self.get_kb(kb_id) + if not kb_helper: + return False + + await kb_helper.delete_vec_db() + async with self.kb_db.get_db() as session: + await session.delete(kb_helper.kb) + await session.commit() + + self.kb_insts.pop(kb_id, None) + return True + + async def list_kbs(self) -> list[KnowledgeBase]: + """列出所有知识库实例""" + kbs = [kb_helper.kb for kb_helper in self.kb_insts.values()] + return kbs + + async def update_kb( + self, + kb_id: str, + kb_name: str, + description: str | None = None, + emoji: str | None = None, + embedding_provider_id: str | None = None, + rerank_provider_id: str | None = None, + chunk_size: int | None = None, + chunk_overlap: int | None = None, + top_k_dense: int | None = None, + top_k_sparse: int | None = None, + top_m_final: int | None = None, + ) -> KBHelper | None: + """更新知识库实例""" + kb_helper = await self.get_kb(kb_id) + if not kb_helper: + return None + + kb = kb_helper.kb + if kb_name is not None: + kb.kb_name = kb_name + if description is not None: + kb.description = description + if emoji is not None: + kb.emoji = emoji + if embedding_provider_id is not None: + kb.embedding_provider_id = embedding_provider_id + kb.rerank_provider_id = rerank_provider_id # 允许设置为 None + if chunk_size is not None: + kb.chunk_size = chunk_size + if chunk_overlap is not None: + kb.chunk_overlap = chunk_overlap + if top_k_dense is not None: + kb.top_k_dense = top_k_dense + if top_k_sparse is not None: + kb.top_k_sparse = top_k_sparse + if top_m_final is not None: + kb.top_m_final = top_m_final + async with self.kb_db.get_db() as session: + session.add(kb) + await session.commit() + await session.refresh(kb) + + return kb_helper + + async def retrieve( + self, + query: str, + kb_names: list[str], + top_k_fusion: int = 20, + top_m_final: int = 5, + ) -> dict | None: + """从指定知识库中检索相关内容""" + kb_ids = [] + kb_id_helper_map = {} + for kb_name in kb_names: + if kb_helper := await self.get_kb_by_name(kb_name): + kb_ids.append(kb_helper.kb.kb_id) + kb_id_helper_map[kb_helper.kb.kb_id] = kb_helper + + if not kb_ids: + return {} + + results = await self.retrieval_manager.retrieve( + query=query, + kb_ids=kb_ids, + kb_id_helper_map=kb_id_helper_map, + top_k_fusion=top_k_fusion, + top_m_final=top_m_final, + ) + if not results: + return None + + context_text = self._format_context(results) + + results_dict = [ + { + "chunk_id": r.chunk_id, + "doc_id": r.doc_id, + "kb_id": r.kb_id, + "kb_name": r.kb_name, + "doc_name": r.doc_name, + "chunk_index": r.metadata.get("chunk_index", 0), + "content": r.content, + "score": r.score, + "char_count": r.metadata.get("char_count", 0), + } + for r in results + ] + + return { + "context_text": context_text, + "results": results_dict, + } + + def _format_context(self, results: list[RetrievalResult]) -> str: + """格式化知识上下文 + + Args: + results: 检索结果列表 + + Returns: + str: 格式化的上下文文本 + """ + lines = ["以下是相关的知识库内容,请参考这些信息回答用户的问题:\n"] + + for i, result in enumerate(results, 1): + lines.append(f"【知识 {i}】") + lines.append(f"来源: {result.kb_name} / {result.doc_name}") + lines.append(f"内容: {result.content}") + lines.append(f"相关度: {result.score:.2f}") + lines.append("") + + return "\n".join(lines) diff --git a/astrbot/core/knowledge_base/kb_sqlite.py b/astrbot/core/knowledge_base/kb_sqlite.py deleted file mode 100644 index 526b6277..00000000 --- a/astrbot/core/knowledge_base/kb_sqlite.py +++ /dev/null @@ -1,230 +0,0 @@ -""" -知识库独立 SQLite 数据库 - -该模块提供知识库专用的独立 SQLite 数据库,与主数据库 (astrbot.db) 完全隔离。 -职责: -- 管理知识库相关表 (knowledge_bases, kb_documents, kb_chunks, kb_media) -- 提供数据库连接和会话管理 -- 执行数据库迁移和初始化 -""" - -from contextlib import asynccontextmanager -from pathlib import Path - -from sqlalchemy import text -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine - -from astrbot.core import logger - - -class KBSQLiteDatabase: - """知识库独立 SQLite 数据库 - - 与主数据库 (astrbot.db) 完全隔离的独立数据库,专门用于存储知识库数据。 - - 特点: - - 数据隔离: 知识库数据不会影响主数据库格式 - - 独立备份: 可以单独备份和恢复知识库数据 - - 性能隔离: 大量知识库查询不会影响主业务性能 - """ - - def __init__(self, db_path: str = "data/knowledge_base/kb.db") -> None: - """初始化知识库数据库 - - Args: - db_path: 数据库文件路径,默认为 data/knowledge_base/kb.db - """ - self.db_path = db_path - self.DATABASE_URL = f"sqlite+aiosqlite:///{db_path}" - self.inited = False - - # 确保目录存在 - Path(db_path).parent.mkdir(parents=True, exist_ok=True) - - # 创建异步引擎 - self.engine = create_async_engine( - self.DATABASE_URL, - echo=False, - pool_pre_ping=True, - pool_recycle=3600, - ) - - # 创建会话工厂 - self.async_session = async_sessionmaker( - self.engine, - class_=AsyncSession, - expire_on_commit=False, - ) - - @asynccontextmanager - async def get_db(self): - """获取数据库会话 - - 用法: - async with kb_db.get_db() as session: - # 执行数据库操作 - result = await session.execute(stmt) - """ - async with self.async_session() as session: - yield session - - async def initialize(self) -> None: - """初始化数据库,创建表并配置 SQLite 参数""" - # noqa: F401 - 这些导入是必需的,用于触发 SQLModel 创建对应的数据库表 - from astrbot.core.knowledge_base.models import ( # noqa: F401 - KBChunk, - KBDocument, - KBMedia, - KBSessionConfig, - KnowledgeBase, - ) - from sqlmodel import SQLModel - - async with self.engine.begin() as conn: - # 创建所有知识库相关表 - await conn.run_sync(SQLModel.metadata.create_all) - - # 配置 SQLite 性能优化参数 - await conn.execute(text("PRAGMA journal_mode=WAL")) - await conn.execute(text("PRAGMA synchronous=NORMAL")) - await conn.execute(text("PRAGMA cache_size=20000")) - await conn.execute(text("PRAGMA temp_store=MEMORY")) - await conn.execute(text("PRAGMA mmap_size=134217728")) - await conn.execute(text("PRAGMA optimize")) - await conn.commit() - - self.inited = True - logger.info(f"知识库数据库已初始化: {self.db_path}") - - async def migrate_to_v1(self) -> None: - """执行知识库数据库 v1 迁移 - - 创建所有必要的索引以优化查询性能 - """ - async with self.get_db() as session: - session: AsyncSession - async with session.begin(): - # 创建知识库表索引 - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_kb_kb_id " - "ON knowledge_bases(kb_id)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_kb_name " - "ON knowledge_bases(kb_name)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_kb_created_at " - "ON knowledge_bases(created_at)" - ) - ) - - # 创建文档表索引 - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_doc_doc_id " - "ON kb_documents(doc_id)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_doc_kb_id " - "ON kb_documents(kb_id)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_doc_name " - "ON kb_documents(doc_name)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_doc_type " - "ON kb_documents(file_type)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_doc_created_at " - "ON kb_documents(created_at)" - ) - ) - - # 创建块表索引 - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_chunk_chunk_id " - "ON kb_chunks(chunk_id)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_chunk_doc_id " - "ON kb_chunks(doc_id)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_chunk_kb_id ON kb_chunks(kb_id)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_chunk_vec_doc_id " - "ON kb_chunks(vec_doc_id)" - ) - ) - - # 创建多媒体表索引 - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_media_media_id " - "ON kb_media(media_id)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_media_doc_id " - "ON kb_media(doc_id)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_media_kb_id ON kb_media(kb_id)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_media_type " - "ON kb_media(media_type)" - ) - ) - - # 创建会话配置表索引 - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_session_config_scope_id " - "ON kb_session_config(scope_id)" - ) - ) - await session.execute( - text( - "CREATE INDEX IF NOT EXISTS idx_session_config_scope " - "ON kb_session_config(scope)" - ) - ) - - await session.commit() - - logger.info("知识库数据库迁移 v1 完成") - - async def close(self) -> None: - """关闭数据库连接""" - await self.engine.dispose() - logger.info(f"知识库数据库已关闭: {self.db_path}") diff --git a/astrbot/core/knowledge_base/manager.py b/astrbot/core/knowledge_base/manager.py deleted file mode 100644 index c4d31ba4..00000000 --- a/astrbot/core/knowledge_base/manager.py +++ /dev/null @@ -1,375 +0,0 @@ -"""知识库管理器 - -该模块提供知识库的CRUD操作和文档上传处理流程。 -""" - -import uuid -from pathlib import Path -from typing import Optional - -import aiofiles -from sqlalchemy import func, select, update - -from astrbot.core.db import BaseDatabase -from astrbot.core.db.vec_db.base import BaseVecDB -from astrbot.core.knowledge_base.chunking.base import BaseChunker -from astrbot.core.knowledge_base.models import KBChunk, KBDocument, KnowledgeBase -from astrbot.core.knowledge_base.parsers.base import BaseParser - - -class KBManager: - """知识库管理器 - - 职责: - - 知识库的 CRUD 操作 - - 文档上传与解析 - - 文档块生成与存储 - - 多媒体资源管理 - """ - - def __init__( - self, - db: BaseDatabase, - vec_db: BaseVecDB, - storage_path: str, - parsers: dict[str, BaseParser], - chunker: BaseChunker, - provider_manager=None, - ): - self.db = db - self.vec_db = vec_db - self.storage_path = Path(storage_path) - self.media_path = self.storage_path / "media" - self.files_path = self.storage_path / "files" - self.parsers = parsers - self.chunker = chunker - self.provider_manager = provider_manager - - # 确保目录存在 - self.media_path.mkdir(parents=True, exist_ok=True) - self.files_path.mkdir(parents=True, exist_ok=True) - - # ===== 知识库操作 ===== - - async def create_kb( - self, - kb_name: str, - description: Optional[str] = None, - emoji: Optional[str] = None, - embedding_provider_id: Optional[str] = None, - rerank_provider_id: Optional[str] = None, - chunk_size: Optional[int] = None, - chunk_overlap: Optional[int] = None, - top_k_dense: Optional[int] = None, - top_k_sparse: Optional[int] = None, - top_m_final: Optional[int] = None, - enable_rerank: Optional[bool] = None, - ) -> KnowledgeBase: - """创建知识库 - - Args: - enable_rerank: 是否启用重排序。 - - 如果明确传入 True/False，则使用该值 - - 如果为 None，则根据是否有可用的 rerank provider 自动决定 - """ - # 智能决定 enable_rerank 的默认值 - if enable_rerank is None: - # 检查是否有可用的 rerank provider - has_rerank_provider = ( - self.provider_manager - and hasattr(self.provider_manager, 'rerank_provider_insts') - and len(self.provider_manager.rerank_provider_insts) > 0 - ) - enable_rerank = has_rerank_provider - - kb = KnowledgeBase( - kb_name=kb_name, - description=description, - emoji=emoji or "📚", - embedding_provider_id=embedding_provider_id, - rerank_provider_id=rerank_provider_id, - chunk_size=chunk_size if chunk_size is not None else 512, - chunk_overlap=chunk_overlap if chunk_overlap is not None else 50, - top_k_dense=top_k_dense if top_k_dense is not None else 50, - top_k_sparse=top_k_sparse if top_k_sparse is not None else 50, - top_m_final=top_m_final if top_m_final is not None else 5, - enable_rerank=enable_rerank, - ) - async with self.db.get_db() as session: - session.add(kb) - await session.commit() - await session.refresh(kb) - return kb - - async def get_kb(self, kb_id: str) -> Optional[KnowledgeBase]: - """获取知识库""" - async with self.db.get_db() as session: - stmt = select(KnowledgeBase).where(KnowledgeBase.kb_id == kb_id) - result = await session.execute(stmt) - return result.scalar_one_or_none() - - async def list_kbs(self, offset: int = 0, limit: int = 100) -> list[KnowledgeBase]: - """列出所有知识库""" - async with self.db.get_db() as session: - stmt = ( - select(KnowledgeBase) - .offset(offset) - .limit(limit) - .order_by(KnowledgeBase.created_at.desc()) - ) - result = await session.execute(stmt) - return list(result.scalars().all()) - - async def update_kb( - self, - kb_id: str, - kb_name: Optional[str] = None, - description: Optional[str] = None, - emoji: Optional[str] = None, - embedding_provider_id: Optional[str] = None, - rerank_provider_id: Optional[str] = None, - chunk_size: Optional[int] = None, - chunk_overlap: Optional[int] = None, - top_k_dense: Optional[int] = None, - top_k_sparse: Optional[int] = None, - top_m_final: Optional[int] = None, - enable_rerank: Optional[bool] = None, - ) -> Optional[KnowledgeBase]: - """更新知识库""" - async with self.db.get_db() as session: - stmt = select(KnowledgeBase).where(KnowledgeBase.kb_id == kb_id) - result = await session.execute(stmt) - kb = result.scalar_one_or_none() - if not kb: - return None - - if kb_name is not None: - kb.kb_name = kb_name - if description is not None: - kb.description = description - if emoji is not None: - kb.emoji = emoji - if embedding_provider_id is not None: - kb.embedding_provider_id = embedding_provider_id - if rerank_provider_id is not None: - kb.rerank_provider_id = rerank_provider_id - if chunk_size is not None: - kb.chunk_size = chunk_size - if chunk_overlap is not None: - kb.chunk_overlap = chunk_overlap - if top_k_dense is not None: - kb.top_k_dense = top_k_dense - if top_k_sparse is not None: - kb.top_k_sparse = top_k_sparse - if top_m_final is not None: - kb.top_m_final = top_m_final - if enable_rerank is not None: - kb.enable_rerank = enable_rerank - - await session.commit() - await session.refresh(kb) - return kb - - async def delete_kb(self, kb_id: str) -> bool: - """删除知识库(级联删除所有文档和资源)""" - # 1. 获取所有文档 - from astrbot.core.knowledge_base.manager_ops import KBManagerOps - - ops = KBManagerOps(self) - docs = await ops.list_documents(kb_id) - - # 2. 删除所有文档(包括文件和向量) - for doc in docs: - await ops.delete_document(doc.doc_id) - - # 3. 删除知识库记录 - async with self.db.get_db() as session: - stmt = select(KnowledgeBase).where(KnowledgeBase.kb_id == kb_id) - result = await session.execute(stmt) - kb = result.scalar_one_or_none() - if not kb: - return False - - await session.delete(kb) - await session.commit() - - return True - - # ===== 文档上传 ===== - - async def upload_document( - self, - kb_id: str, - file_name: str, - file_content: bytes, - file_type: str, - ) -> KBDocument: - """上传并处理文档（带原子性保证和失败清理） - - 流程: - 1. 保存原始文件 - 2. 解析文档内容 - 3. 提取多媒体资源 - 4. 分块处理 - 5. 生成向量并存储 - 6. 保存元数据（事务） - 7. 更新统计 - """ - doc_id = str(uuid.uuid4()) - file_path = None - media_paths = [] - vec_doc_ids = [] - - try: - # 1. 保存原始文件 - file_path = self.files_path / kb_id / f"{doc_id}.{file_type}" - file_path.parent.mkdir(parents=True, exist_ok=True) - - async with aiofiles.open(file_path, "wb") as f: - await f.write(file_content) - - # 2. 解析文档 - parser = self.parsers.get(file_type) - if not parser: - raise ValueError(f"不支持的文件类型: {file_type}") - - parse_result = await parser.parse(file_content, file_name) - text_content = parse_result.text - media_items = parse_result.media - - # 3. 保存多媒体资源 - from astrbot.core.knowledge_base.manager_ops import KBManagerOps - - ops = KBManagerOps(self) - saved_media = [] - for media_item in media_items: - media = await ops._save_media( - kb_id=kb_id, - doc_id=doc_id, - media_type=media_item.media_type, - file_name=media_item.file_name, - content=media_item.content, - mime_type=media_item.mime_type, - ) - saved_media.append(media) - media_paths.append(Path(media.file_path)) - - # 4. 文档分块 - chunks_text = await self.chunker.chunk(text_content) - - # 5. 生成向量并存储 - saved_chunks = [] - for idx, chunk_text in enumerate(chunks_text): - # 存储到向量数据库 - vec_doc_id = await self.vec_db.insert( - content=chunk_text, - metadata={ - "kb_id": kb_id, - "doc_id": doc_id, - "chunk_index": idx, - }, - ) - vec_doc_ids.append(str(vec_doc_id)) - - # 保存块元数据 - chunk = KBChunk( - doc_id=doc_id, - kb_id=kb_id, - chunk_index=idx, - content=chunk_text, - char_count=len(chunk_text), - vec_doc_id=str(vec_doc_id), - ) - saved_chunks.append(chunk) - - # 6. 保存文档元数据（事务） - doc = KBDocument( - doc_id=doc_id, - kb_id=kb_id, - doc_name=file_name, - file_type=file_type, - file_size=len(file_content), - file_path=str(file_path), - chunk_count=len(saved_chunks), - media_count=len(saved_media), - ) - - async with self.db.get_db() as session: - async with session.begin(): - session.add(doc) - for chunk in saved_chunks: - session.add(chunk) - for media in saved_media: - session.add(media) - await session.commit() - - await session.refresh(doc) - - # 7. 更新知识库统计 - await self._update_kb_stats(kb_id) - - return doc - - except Exception as e: - # 失败清理：删除已创建的资源 - from astrbot.core import logger - - logger.error(f"文档上传失败，开始清理资源: {e}") - - # 清理向量数据库 - for vec_id in vec_doc_ids: - try: - await self.vec_db.delete(vec_id) - except Exception as ve: - logger.warning(f"清理向量失败 {vec_id}: {ve}") - - # 清理多媒体文件 - for media_path in media_paths: - try: - if media_path.exists(): - media_path.unlink() - except Exception as me: - logger.warning(f"清理多媒体文件失败 {media_path}: {me}") - - # 清理文档文件 - if file_path and file_path.exists(): - try: - file_path.unlink() - except Exception as fe: - logger.warning(f"清理文档文件失败 {file_path}: {fe}") - - # 重新抛出原始异常 - raise - - # ===== 统计更新 ===== - - async def _update_kb_stats(self, kb_id: str): - """更新知识库统计信息（事务中执行）""" - async with self.db.get_db() as session: - async with session.begin(): - # 统计文档数（在事务中查询） - doc_count = ( - await session.scalar( - select(func.count(KBDocument.id)).where( - KBDocument.kb_id == kb_id - ) - ) - or 0 - ) - - # 统计块数（在事务中查询） - chunk_count = ( - await session.scalar( - select(func.count(KBChunk.id)).where(KBChunk.kb_id == kb_id) - ) - or 0 - ) - - # 更新知识库（在同一事务中） - await session.execute( - update(KnowledgeBase) - .where(KnowledgeBase.kb_id == kb_id) - .values(doc_count=doc_count, chunk_count=chunk_count) - ) - - await session.commit() diff --git a/astrbot/core/knowledge_base/manager_ops.py b/astrbot/core/knowledge_base/manager_ops.py deleted file mode 100644 index e0ab5f6d..00000000 --- a/astrbot/core/knowledge_base/manager_ops.py +++ /dev/null @@ -1,312 +0,0 @@ -"""知识库管理器辅助操作 - -该模块提供文档、块和多媒体的管理操作。 -""" - -import uuid -from pathlib import Path -from typing import TYPE_CHECKING - -import aiofiles -from sqlalchemy import delete, func, select - -from astrbot.core.knowledge_base.models import KBChunk, KBDocument, KBMedia - -if TYPE_CHECKING: - from astrbot.core.knowledge_base.manager import KBManager - - -class KBManagerOps: - """知识库管理器辅助操作类 - - 职责: - - 文档管理操作 - - 块管理操作 - - 多媒体管理操作 - """ - - def __init__(self, manager: "KBManager"): - self.manager = manager - self.db = manager.db - self.vec_db = manager.vec_db - self.media_path = manager.media_path - self.files_path = manager.files_path - - # ===== 文档操作 ===== - - async def list_documents( - self, kb_id: str, offset: int = 0, limit: int = 100 - ) -> list[KBDocument]: - """列出知识库的所有文档""" - async with self.db.get_db() as session: - stmt = ( - select(KBDocument) - .where(KBDocument.kb_id == kb_id) - .offset(offset) - .limit(limit) - .order_by(KBDocument.created_at.desc()) - ) - result = await session.execute(stmt) - return list(result.scalars().all()) - - async def get_document(self, doc_id: str) -> KBDocument | None: - """获取文档详情""" - async with self.db.get_db() as session: - stmt = select(KBDocument).where(KBDocument.doc_id == doc_id) - result = await session.execute(stmt) - return result.scalar_one_or_none() - - async def delete_document(self, doc_id: str) -> bool: - """删除文档(级联删除块、多媒体、向量) - - 采用三阶段删除策略: - 1. 删除向量数据库中的向量(允许部分失败) - 2. 删除SQL数据库中的记录(事务保证原子性) - 3. 删除文件系统中的文件(失败不影响数据一致性) - """ - from astrbot.core import logger - - # 0. 获取文档信息 - doc = await self.get_document(doc_id) - if not doc: - return False - - # 收集所有需要删除的资源 - chunks = await self.list_chunks(doc_id) - media_list = await self.list_media(doc_id) - - # ===== 第一阶段: 删除向量(可重试) ===== - vec_ids_to_delete = [chunk.vec_doc_id for chunk in chunks] - deleted_vec_ids = [] - failed_vec_ids = [] - - for vec_id in vec_ids_to_delete: - try: - await self.vec_db.delete(vec_id) - deleted_vec_ids.append(vec_id) - except Exception as e: - logger.error(f"删除向量失败: {vec_id}, {e}") - failed_vec_ids.append(vec_id) - - # 如果向量删除失败过多(超过50%),中止操作 - if len(failed_vec_ids) > len(vec_ids_to_delete) * 0.5: - logger.error( - f"向量删除失败过多 ({len(failed_vec_ids)}/{len(vec_ids_to_delete)}), 中止文档删除" - ) - return False - - # 记录部分失败但继续执行 - if failed_vec_ids: - logger.warning( - f"部分向量删除失败 ({len(failed_vec_ids)}/{len(vec_ids_to_delete)}), 但继续执行删除操作" - ) - - # ===== 第二阶段: 删除数据库记录(事务) ===== - async with self.db.get_db() as session: - async with session.begin(): - # 删除块记录 - await session.execute(delete(KBChunk).where(KBChunk.doc_id == doc_id)) - - # 删除多媒体记录 - await session.execute(delete(KBMedia).where(KBMedia.doc_id == doc_id)) - - # 删除文档记录 - await session.execute( - delete(KBDocument).where(KBDocument.doc_id == doc_id) - ) - - await session.commit() - - # ===== 第三阶段: 删除文件(失败不影响) ===== - # 删除多媒体文件 - for media in media_list: - try: - media_path = Path(media.file_path) - if media_path.exists(): - media_path.unlink() - except Exception as e: - logger.warning(f"删除多媒体文件失败: {media.file_path}, {e}") - - # 删除文档文件 - try: - file_path = Path(doc.file_path) - if file_path.exists(): - file_path.unlink() - except Exception as e: - logger.warning(f"删除文档文件失败: {doc.file_path}, {e}") - - # ===== 更新统计 ===== - await self.manager._update_kb_stats(doc.kb_id) - - return True - - # ===== 块操作 ===== - - async def list_chunks(self, doc_id: str) -> list[KBChunk]: - """列出文档的所有块""" - async with self.db.get_db() as session: - stmt = ( - select(KBChunk) - .where(KBChunk.doc_id == doc_id) - .order_by(KBChunk.chunk_index) - ) - result = await session.execute(stmt) - return list(result.scalars().all()) - - async def delete_chunk(self, chunk_id: str) -> bool: - """删除单个块 - - 流程: - 1. 查询块信息 - 2. 删除向量 - 3. 删除数据库记录 - 4. 更新文档统计 - """ - from astrbot.core import logger - - # 1. 查询块信息 - async with self.db.get_db() as session: - stmt = select(KBChunk).where(KBChunk.chunk_id == chunk_id) - result = await session.execute(stmt) - chunk = result.scalar_one_or_none() - if not chunk: - return False - - doc_id = chunk.doc_id - vec_doc_id = chunk.vec_doc_id - - # 2. 删除向量 - try: - await self.vec_db.delete(vec_doc_id) - except Exception as e: - logger.error(f"删除向量失败: {vec_doc_id}, {e}") - return False - - # 3. 删除数据库记录 - async with self.db.get_db() as session: - async with session.begin(): - await session.execute( - delete(KBChunk).where(KBChunk.chunk_id == chunk_id) - ) - await session.commit() - - # 4. 更新文档统计 - await self._update_doc_stats(doc_id) - - return True - - # ===== 多媒体操作 ===== - - async def list_media(self, doc_id: str) -> list[KBMedia]: - """列出文档的所有多媒体资源""" - async with self.db.get_db() as session: - stmt = select(KBMedia).where(KBMedia.doc_id == doc_id) - result = await session.execute(stmt) - return list(result.scalars().all()) - - async def delete_media(self, media_id: str) -> bool: - """删除多媒体资源 - - 流程: - 1. 查询媒体信息 - 2. 删除数据库记录 - 3. 删除文件(失败不影响) - 4. 更新文档统计 - """ - from astrbot.core import logger - - # 1. 查询媒体信息 - async with self.db.get_db() as session: - stmt = select(KBMedia).where(KBMedia.media_id == media_id) - result = await session.execute(stmt) - media = result.scalar_one_or_none() - if not media: - return False - - doc_id = media.doc_id - file_path_str = media.file_path - - # 2. 删除数据库记录 - async with self.db.get_db() as session: - async with session.begin(): - await session.execute( - delete(KBMedia).where(KBMedia.media_id == media_id) - ) - await session.commit() - - # 3. 删除文件(失败不影响) - try: - media_path = Path(file_path_str) - if media_path.exists(): - media_path.unlink() - except Exception as e: - logger.warning(f"删除多媒体文件失败: {file_path_str}, {e}") - - # 4. 更新文档统计 - await self._update_doc_stats(doc_id) - - return True - - # ===== 内部辅助方法 ===== - - async def _save_media( - self, - kb_id: str, - doc_id: str, - media_type: str, - file_name: str, - content: bytes, - mime_type: str, - ) -> KBMedia: - """保存多媒体资源""" - media_id = str(uuid.uuid4()) - ext = Path(file_name).suffix - - # 保存文件 - file_path = self.media_path / kb_id / doc_id / f"{media_id}{ext}" - file_path.parent.mkdir(parents=True, exist_ok=True) - - async with aiofiles.open(file_path, "wb") as f: - await f.write(content) - - # 创建记录 - media = KBMedia( - media_id=media_id, - doc_id=doc_id, - kb_id=kb_id, - media_type=media_type, - file_name=file_name, - file_path=str(file_path), - file_size=len(content), - mime_type=mime_type, - ) - - return media - - async def _update_doc_stats(self, doc_id: str): - """更新文档统计信息(事务中执行)""" - async with self.db.get_db() as session: - async with session.begin(): - # 统计块数 - chunk_count = ( - await session.scalar( - select(func.count(KBChunk.id)).where(KBChunk.doc_id == doc_id) - ) - ) or 0 - - # 统计多媒体数 - media_count = ( - await session.scalar( - select(func.count(KBMedia.id)).where(KBMedia.doc_id == doc_id) - ) - ) or 0 - - # 更新文档 - doc = await session.scalar( - select(KBDocument).where(KBDocument.doc_id == doc_id) - ) - if doc: - doc.chunk_count = chunk_count - doc.media_count = media_count - - await session.commit() diff --git a/astrbot/core/knowledge_base/models.py b/astrbot/core/knowledge_base/models.py index 6ec77ad9..010d6113 100644 --- a/astrbot/core/knowledge_base/models.py +++ b/astrbot/core/knowledge_base/models.py @@ -1,31 +1,20 @@ -"""知识库管理功能的数据模型定义 - -该模块定义了知识库系统所需的数据模型,包括: -- KnowledgeBase: 知识库表 (存储在独立的 kb.db) -- KBDocument: 文档表 (存储在独立的 kb.db) -- KBChunk: 文档块表 (存储在独立的 kb.db) -- KBMedia: 多媒体资源表 (存储在独立的 kb.db) -- KBSessionConfig: 会话配置表 (存储在独立的 kb.db) - -注意: -- 所有模型存储在独立的知识库数据库 (data/knowledge_base/kb.db) -- 与主数据库 (astrbot.db) 完全解耦 -""" - import uuid from datetime import datetime, timezone -from typing import Optional -from sqlmodel import Field, SQLModel, Text, UniqueConstraint +from sqlmodel import Field, SQLModel, Text, UniqueConstraint, MetaData -class KnowledgeBase(SQLModel, table=True): +class BaseKBModel(SQLModel, table=False): + metadata = MetaData() + + +class KnowledgeBase(BaseKBModel, table=True): """知识库表 存储知识库的基本信息和统计数据。 """ - __tablename__ = "knowledge_bases" + __tablename__ = "knowledge_bases" # type: ignore id: int | None = Field( primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None @@ -38,18 +27,17 @@ class KnowledgeBase(SQLModel, table=True): index=True, ) kb_name: str = Field(max_length=100, nullable=False) - description: Optional[str] = Field(default=None, sa_type=Text) - emoji: Optional[str] = Field(default="📚", max_length=10) - embedding_provider_id: Optional[str] = Field(default=None, max_length=100) - rerank_provider_id: Optional[str] = Field(default=None, max_length=100) + description: str | None = Field(default=None, sa_type=Text) + emoji: str | None = Field(default="📚", max_length=10) + embedding_provider_id: str | None = Field(default=None, max_length=100) + rerank_provider_id: str | None = Field(default=None, max_length=100) # 分块配置参数 - chunk_size: Optional[int] = Field(default=512, nullable=True) - chunk_overlap: Optional[int] = Field(default=50, nullable=True) + chunk_size: int | None = Field(default=512, nullable=True) + chunk_overlap: int | None = Field(default=50, nullable=True) # 检索配置参数 - top_k_dense: Optional[int] = Field(default=50, nullable=True) - top_k_sparse: Optional[int] = Field(default=50, nullable=True) - top_m_final: Optional[int] = Field(default=5, nullable=True) - enable_rerank: Optional[bool] = Field(default=False, nullable=True) + top_k_dense: int | None = Field(default=50, nullable=True) + top_k_sparse: int | None = Field(default=50, nullable=True) + top_m_final: int | None = Field(default=5, nullable=True) created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) updated_at: datetime = Field( default_factory=lambda: datetime.now(timezone.utc), @@ -58,14 +46,21 @@ class KnowledgeBase(SQLModel, table=True): doc_count: int = Field(default=0, nullable=False) chunk_count: int = Field(default=0, nullable=False) + __table_args__ = ( + UniqueConstraint( + "kb_name", + name="uix_kb_name", + ), + ) -class KBDocument(SQLModel, table=True): + +class KBDocument(BaseKBModel, table=True): """文档表 存储上传到知识库的文档元数据。 """ - __tablename__ = "kb_documents" + __tablename__ = "kb_documents" # type: ignore id: int | None = Field( primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None @@ -91,40 +86,13 @@ class KBDocument(SQLModel, table=True): ) -class KBChunk(SQLModel, table=True): - """文档块表 - - 存储文档分块后的文本内容和向量索引关联信息。 - """ - - __tablename__ = "kb_chunks" - - id: int | None = Field( - primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None - ) - chunk_id: str = Field( - max_length=36, - nullable=False, - unique=True, - default_factory=lambda: str(uuid.uuid4()), - index=True, - ) - doc_id: str = Field(max_length=36, nullable=False, index=True) - kb_id: str = Field(max_length=36, nullable=False, index=True) - chunk_index: int = Field(nullable=False) - content: str = Field(sa_type=Text, nullable=False) - char_count: int = Field(nullable=False) - vec_doc_id: str = Field(max_length=100, nullable=False, index=True) - created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) - - -class KBMedia(SQLModel, table=True): +class KBMedia(BaseKBModel, table=True): """多媒体资源表 存储从文档中提取的图片、视频等多媒体资源。 """ - __tablename__ = "kb_media" + __tablename__ = "kb_media" # type: ignore id: int | None = Field( primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None @@ -144,39 +112,3 @@ class KBMedia(SQLModel, table=True): file_size: int = Field(nullable=False) mime_type: str = Field(max_length=100, nullable=False) created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) - - -class KBSessionConfig(SQLModel, table=True): - """会话知识库配置表 - - 存储会话或平台级别的知识库关联配置。 - 该表存储在知识库独立数据库中,保持完全解耦。 - - 支持两种配置范围: - - platform: 平台级别配置 (如 'qq', 'telegram') - - session: 会话级别配置 (如 'qq:group:12345') - """ - - __tablename__ = "kb_session_config" - - id: int | None = Field( - primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None - ) - config_id: str = Field( - max_length=36, - nullable=False, - unique=True, - default_factory=lambda: str(uuid.uuid4()), - ) - scope: str = Field(max_length=20, nullable=False) - scope_id: str = Field(max_length=255, nullable=False, index=True) - kb_ids: str = Field(sa_type=Text, nullable=False) - top_k: Optional[int] = Field(default=None, nullable=True) - enable_rerank: Optional[bool] = Field(default=None, nullable=True) - created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) - updated_at: datetime = Field( - default_factory=lambda: datetime.now(timezone.utc), - sa_column_kwargs={"onupdate": datetime.now(timezone.utc)}, - ) - - __table_args__ = (UniqueConstraint("scope", "scope_id", name="uix_scope_scope_id"),) diff --git a/astrbot/core/knowledge_base/parsers/pdf_parser.py b/astrbot/core/knowledge_base/parsers/pdf_parser.py index 8bb1dea6..fca62687 100644 --- a/astrbot/core/knowledge_base/parsers/pdf_parser.py +++ b/astrbot/core/knowledge_base/parsers/pdf_parser.py @@ -51,10 +51,10 @@ class PDFParser(BaseParser): continue resources = page["/Resources"] - if not resources or "/XObject" not in resources: + if not resources or "/XObject" not in resources: # type: ignore continue - xobjects = resources["/XObject"].get_object() + xobjects = resources["/XObject"].get_object() # type: ignore if not xobjects: continue diff --git a/astrbot/core/knowledge_base/retrieval/hit_stopwords.txt b/astrbot/core/knowledge_base/retrieval/hit_stopwords.txt new file mode 100644 index 00000000..84b26283 --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/hit_stopwords.txt @@ -0,0 +1,767 @@ +——— +》）， +）÷（１－ +”， +）、 +＝（ +: +→ +℃ +& +* +一一 +~~~~ +’ +. +『 +.一 +./ +-- +』 +＝″ +【 +［＊］ +｝＞ +［⑤］］ +［①Ｄ］ +ｃ］ +ｎｇ昉 +＊ +// +［ +］ +［②ｅ］ +［②ｇ］ +＝｛ +} +，也 +‘ +Ａ +［①⑥］ +［②Ｂ］ +［①ａ］ +［④ａ］ +［①③］ +［③ｈ］ +③］ +１． +－－ +［②ｂ］ +’‘ +××× +［①⑧］ +０：２ +＝［ +［⑤ｂ］ +［②ｃ］ +［④ｂ］ +［②③］ +［③ａ］ +［④ｃ］ +［①⑤］ +［①⑦］ +［①ｇ］ +∈［ +［①⑨］ +［①④］ +［①ｃ］ +［②ｆ］ +［②⑧］ +［②①］ +［①Ｃ］ +［③ｃ］ +［③ｇ］ +［②⑤］ +［②②］ +一. +［①ｈ］ +.数 +［］ +［①Ｂ］ +数/ +［①ｉ］ +［③ｅ］ +［①①］ +［④ｄ］ +［④ｅ］ +［③ｂ］ +［⑤ａ］ +［①Ａ］ +［②⑧］ +［②⑦］ +［①ｄ］ +［②ｊ］ +〕〔 +］［ +:// +′∈ +［②④ +［⑤ｅ］ +１２％ +ｂ］ +... +................... +…………………………………………………③ +ＺＸＦＩＴＬ +［③Ｆ］ +」 +［①ｏ］ +］∧′＝［ +∪φ∈ +′｜ +｛－ +②ｃ +｝ +［③①］ +Ｒ．Ｌ． +［①Ｅ］ +Ψ +－［＊］－ +↑ +.日 +［②ｄ］ +［② +［②⑦］ +［②②］ +［③ｅ］ +［①ｉ］ +［①Ｂ］ +［①ｈ］ +［①ｄ］ +［①ｇ］ +［①②］ +［②ａ］ +ｆ］ +［⑩］ +ａ］ +［①ｅ］ +［②ｈ］ +［②⑥］ +［③ｄ］ +［②⑩］ +ｅ］ +〉 +】 +元／吨 +［②⑩］ +２．３％ +５：０ +［①］ +:: +［②］ +［③］ +［④］ +［⑤］ +［⑥］ +［⑦］ +［⑧］ +［⑨］ +…… +—— +? +、 +。 +“ +” +《 +》 +！ +， +： +； +？ +． +, +． +' +? +· +——— +── +? +— +< +> +（ +） +〔 +〕 +[ +] +( +) +- ++ +～ +× +／ +/ +① +② +③ +④ +⑤ +⑥ +⑦ +⑧ +⑨ +⑩ +Ⅲ +В +" +; +# +@ +γ +μ +φ +φ． +× +Δ +■ +▲ +sub +exp +sup +sub +Lex +＃ +％ +＆ +＇ +＋ +＋ξ +＋＋ +－ +－β +＜ +＜± +＜Δ +＜λ +＜φ +＜＜ += +＝ +＝☆ +＝－ +＞ +＞λ +＿ +～± +～＋ +［⑤ｆ］ +［⑤ｄ］ +［②ｉ］ +≈ +［②Ｇ］ +［①ｆ］ +ＬＩ +㈧ +［－ +...... +〉 +［③⑩］ +第二 +一番 +一直 +一个 +一些 +许多 +种 +有的是 +也就是说 +末##末 +啊 +阿 +哎 +哎呀 +哎哟 +唉 +俺 +俺们 +按 +按照 +吧 +吧哒 +把 +罢了 +被 +本 +本着 +比 +比方 +比如 +鄙人 +彼 +彼此 +边 +别 +别的 +别说 +并 +并且 +不比 +不成 +不单 +不但 +不独 +不管 +不光 +不过 +不仅 +不拘 +不论 +不怕 +不然 +不如 +不特 +不惟 +不问 +不只 +朝 +朝着 +趁 +趁着 +乘 +冲 +除 +除此之外 +除非 +除了 +此 +此间 +此外 +从 +从而 +打 +待 +但 +但是 +当 +当着 +到 +得 +的 +的话 +等 +等等 +地 +第 +叮咚 +对 +对于 +多 +多少 +而 +而况 +而且 +而是 +而外 +而言 +而已 +尔后 +反过来 +反过来说 +反之 +非但 +非徒 +否则 +嘎 +嘎登 +该 +赶 +个 +各 +各个 +各位 +各种 +各自 +给 +根据 +跟 +故 +故此 +固然 +关于 +管 +归 +果然 +果真 +过 +哈 +哈哈 +呵 +和 +何 +何处 +何况 +何时 +嘿 +哼 +哼唷 +呼哧 +乎 +哗 +还是 +还有 +换句话说 +换言之 +或 +或是 +或者 +极了 +及 +及其 +及至 +即 +即便 +即或 +即令 +即若 +即使 +几 +几时 +己 +既 +既然 +既是 +继而 +加之 +假如 +假若 +假使 +鉴于 +将 +较 +较之 +叫 +接着 +结果 +借 +紧接着 +进而 +尽 +尽管 +经 +经过 +就 +就是 +就是说 +据 +具体地说 +具体说来 +开始 +开外 +靠 +咳 +可 +可见 +可是 +可以 +况且 +啦 +来 +来着 +离 +例如 +哩 +连 +连同 +两者 +了 +临 +另 +另外 +另一方面 +论 +嘛 +吗 +慢说 +漫说 +冒 +么 +每 +每当 +们 +莫若 +某 +某个 +某些 +拿 +哪 +哪边 +哪儿 +哪个 +哪里 +哪年 +哪怕 +哪天 +哪些 +哪样 +那 +那边 +那儿 +那个 +那会儿 +那里 +那么 +那么些 +那么样 +那时 +那些 +那样 +乃 +乃至 +呢 +能 +你 +你们 +您 +宁 +宁可 +宁肯 +宁愿 +哦 +呕 +啪达 +旁人 +呸 +凭 +凭借 +其 +其次 +其二 +其他 +其它 +其一 +其余 +其中 +起 +起见 +起见 +岂但 +恰恰相反 +前后 +前者 +且 +然而 +然后 +然则 +让 +人家 +任 +任何 +任凭 +如 +如此 +如果 +如何 +如其 +如若 +如上所述 +若 +若非 +若是 +啥 +上下 +尚且 +设若 +设使 +甚而 +甚么 +甚至 +省得 +时候 +什么 +什么样 +使得 +是 +是的 +首先 +谁 +谁知 +顺 +顺着 +似的 +虽 +虽然 +虽说 +虽则 +随 +随着 +所 +所以 +他 +他们 +他人 +它 +它们 +她 +她们 +倘 +倘或 +倘然 +倘若 +倘使 +腾 +替 +通过 +同 +同时 +哇 +万一 +往 +望 +为 +为何 +为了 +为什么 +为着 +喂 +嗡嗡 +我 +我们 +呜 +呜呼 +乌乎 +无论 +无宁 +毋宁 +嘻 +吓 +相对而言 +像 +向 +向着 +嘘 +呀 +焉 +沿 +沿着 +要 +要不 +要不然 +要不是 +要么 +要是 +也 +也罢 +也好 +一 +一般 +一旦 +一方面 +一来 +一切 +一样 +一则 +依 +依照 +矣 +以 +以便 +以及 +以免 +以至 +以至于 +以致 +抑或 +因 +因此 +因而 +因为 +哟 +用 +由 +由此可见 +由于 +有 +有的 +有关 +有些 +又 +于 +于是 +于是乎 +与 +与此同时 +与否 +与其 +越是 +云云 +哉 +再说 +再者 +在 +在下 +咱 +咱们 +则 +怎 +怎么 +怎么办 +怎么样 +怎样 +咋 +照 +照着 +者 +这 +这边 +这儿 +这个 +这会儿 +这就是说 +这里 +这么 +这么点儿 +这么些 +这么样 +这时 +这些 +这样 +正如 +吱 +之 +之类 +之所以 +之一 +只是 +只限 +只要 +只有 +至 +至于 +诸位 +着 +着呢 +自 +自从 +自个儿 +自各儿 +自己 +自家 +自身 +综上所述 +总的来看 +总的来说 +总的说来 +总而言之 +总之 +纵 +纵令 +纵然 +纵使 +遵照 +作为 +兮 +呃 +呗 +咚 +咦 +喏 +啐 +喔唷 +嗬 +嗯 +嗳 \ No newline at end of file diff --git a/astrbot/core/knowledge_base/retrieval/manager.py b/astrbot/core/knowledge_base/retrieval/manager.py index fb19d671..278e4da2 100644 --- a/astrbot/core/knowledge_base/retrieval/manager.py +++ b/astrbot/core/knowledge_base/retrieval/manager.py @@ -3,15 +3,19 @@ 协调稠密检索、稀疏检索和 Rerank,提供统一的检索接口 """ -import json -from dataclasses import dataclass -from typing import List, Optional +import time -from astrbot.core.db.vec_db.base import BaseVecDB -from astrbot.core.knowledge_base.database import KBDatabase +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 @@ -38,36 +42,29 @@ class RetrievalManager: def __init__( self, - vec_db: BaseVecDB, sparse_retriever: SparseRetriever, rank_fusion: RankFusion, - kb_db: KBDatabase, - rerank_provider: Optional[RerankProvider] = None, + kb_db: KBSQLiteDatabase, ): """初始化检索管理器 Args: - vec_db: 向量数据库实例 + vec_db_factory: 向量数据库工厂 sparse_retriever: 稀疏检索器 rank_fusion: 结果融合器 kb_db: 知识库数据库实例 - rerank_provider: Rerank 提供商 (可选) """ - self.vec_db = vec_db self.sparse_retriever = sparse_retriever self.rank_fusion = rank_fusion self.kb_db = kb_db - self.rerank_provider = rerank_provider async def retrieve( self, query: str, kb_ids: List[str], - top_k_dense: int = 50, - top_k_sparse: int = 50, - top_n_fusion: int = 20, + kb_id_helper_map: dict[str, KBHelper], + top_k_fusion: int = 20, top_m_final: int = 5, - enable_rerank: bool = True, ) -> List[RetrievalResult]: """混合检索 @@ -80,40 +77,74 @@ class RetrievalManager: Args: query: 查询文本 kb_ids: 知识库 ID 列表 - top_k_dense: 稠密检索返回数量 - top_k_sparse: 稀疏检索返回数量 - top_n_fusion: 融合后返回数量 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, - top_k=top_k_dense, + 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, - top_k=top_k_sparse, + 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_n_fusion, + 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_chunk_with_metadata(fr.chunk_id) + metadata_dict = await self.kb_db.get_document_with_metadata(fr.doc_id) if metadata_dict: retrieval_results.append( RetrievalResult( @@ -125,33 +156,45 @@ class RetrievalManager: content=fr.content, score=fr.score, metadata={ - "chunk_index": metadata_dict["chunk"].chunk_index, - "char_count": metadata_dict["chunk"].char_count, + "chunk_index": fr.chunk_index, + "char_count": len(fr.content), }, ) ) - # 5. Rerank (可选) - if enable_rerank and self.rerank_provider and retrieval_results: + # 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=self.rerank_provider, + rerank_provider=first_rerank, ) - else: - retrieval_results = retrieval_results[:top_m_final] - return retrieval_results + return retrieval_results[:top_m_final] async def _dense_retrieve( self, query: str, kb_ids: List[str], - top_k: int, + kb_options: dict, ): """稠密检索 (向量相似度) + 为每个知识库使用独立的向量数据库进行检索,然后合并结果。 + Args: query: 查询文本 kb_ids: 知识库 ID 列表 @@ -160,28 +203,32 @@ class RetrievalManager: Returns: List[Result]: 检索结果列表 """ - # 直接调用向量数据库检索 - vec_results = await self.vec_db.retrieve( - query=query, - top_k=top_k * len(kb_ids) * 2, # 增加候选数量以便过滤 - ) - - # 过滤:只保留指定知识库的结果 - filtered_results = [] - for result in vec_results: - metadata_str = result.data.get("metadata", "{}") + all_results: list[Result] = [] + for kb_id in kb_ids: + if kb_id not in kb_options: + continue try: - metadata = json.loads(metadata_str) - except (json.JSONDecodeError, TypeError): - metadata = {} + 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}, + ) - if metadata.get("kb_id") in kb_ids: - filtered_results.append(result) + all_results.extend(vec_results) + except Exception as e: + from astrbot.core import logger - if len(filtered_results) >= top_k: - break + logger.warning(f"知识库 {kb_id} 稠密检索失败: {e}") + continue - return filtered_results[:top_k] + # 按相似度排序并返回 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, diff --git a/astrbot/core/knowledge_base/retrieval/rank_fusion.py b/astrbot/core/knowledge_base/retrieval/rank_fusion.py index b05fe1be..3ceba4ff 100644 --- a/astrbot/core/knowledge_base/retrieval/rank_fusion.py +++ b/astrbot/core/knowledge_base/retrieval/rank_fusion.py @@ -3,11 +3,11 @@ 使用 Reciprocal Rank Fusion (RRF) 算法融合稠密检索和稀疏检索的结果 """ +import json from dataclasses import dataclass -from typing import Dict, List from astrbot.core.db.vec_db.base import Result -from astrbot.core.knowledge_base.database import KBDatabase +from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase from astrbot.core.knowledge_base.retrieval.sparse_retriever import SparseResult @@ -16,6 +16,7 @@ class FusedResult: """融合后的检索结果""" chunk_id: str + chunk_index: int doc_id: str kb_id: str content: str @@ -30,7 +31,7 @@ class RankFusion: - 使用 Reciprocal Rank Fusion (RRF) 算法 """ - def __init__(self, kb_db: KBDatabase, k: int = 60): + def __init__(self, kb_db: KBSQLiteDatabase, k: int = 60): """初始化结果融合器 Args: @@ -42,10 +43,10 @@ class RankFusion: async def fuse( self, - dense_results: List[Result], - sparse_results: List[SparseResult], + dense_results: list[Result], + sparse_results: list[SparseResult], top_k: int = 20, - ) -> List[FusedResult]: + ) -> list[FusedResult]: """融合稠密和稀疏检索结果 RRF 公式: @@ -62,14 +63,14 @@ class RankFusion: # 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 (来自稠密检索的是 vec_doc_id, 稀疏检索的是 chunk_id) + # 2. 收集所有唯一的 ID # 需要统一为 chunk_id all_chunk_ids = set() - vec_doc_id_to_dense = {} # vec_doc_id -> Result - chunk_id_to_sparse = {} # chunk_id -> SparseResult + 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: @@ -83,7 +84,7 @@ class RankFusion: vec_doc_id_to_dense[vec_doc_id] = r # 3. 计算 RRF 分数 - rrf_scores: Dict[str, float] = {} + rrf_scores: dict[str, float] = {} for identifier in all_chunk_ids: score = 0.0 @@ -112,6 +113,7 @@ class RankFusion: 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, @@ -120,16 +122,17 @@ class RankFusion: ) elif identifier in vec_doc_id_to_dense: # 从向量检索获取信息,需要从数据库获取块的详细信息 - chunk = await self.kb_db.get_chunk_by_vec_doc_id(identifier) - if chunk: - fused_results.append( - FusedResult( - chunk_id=chunk.chunk_id, - doc_id=chunk.doc_id, - kb_id=chunk.kb_id, - content=chunk.content, - score=rrf_scores[identifier], - ) + 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 diff --git a/astrbot/core/knowledge_base/retrieval/sparse_retriever.py b/astrbot/core/knowledge_base/retrieval/sparse_retriever.py index 15c20512..315930b3 100644 --- a/astrbot/core/knowledge_base/retrieval/sparse_retriever.py +++ b/astrbot/core/knowledge_base/retrieval/sparse_retriever.py @@ -3,18 +3,20 @@ 使用 BM25 算法进行基于关键词的文档检索 """ +import jieba +import os +import json from dataclasses import dataclass -from typing import List - from rank_bm25 import BM25Okapi - -from astrbot.core.knowledge_base.database import KBDatabase +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 @@ -30,7 +32,7 @@ class SparseRetriever: - 使用 BM25 算法计算相关度 """ - def __init__(self, kb_db: KBDatabase): + def __init__(self, kb_db: KBSQLiteDatabase): """初始化稀疏检索器 Args: @@ -39,37 +41,73 @@ class SparseRetriever: 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], - top_k: int = 50, - ) -> List[SparseResult]: + kb_ids: list[str], + kb_options: dict, + ) -> list[SparseResult]: """执行稀疏检索 Args: query: 查询文本 kb_ids: 知识库 ID 列表 - top_k: 返回结果数量 + kb_options: 每个知识库的检索选项 Returns: List[SparseResult]: 检索结果列表 """ # 1. 获取所有相关块 - chunks = await self.kb_db.get_chunks_by_kb_ids(kb_ids) + 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.content for chunk in chunks] - tokenized_corpus = [doc.split() for doc in corpus] + 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 = query.split() + 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 @@ -78,13 +116,15 @@ class SparseRetriever: chunk = chunks[idx] results.append( SparseResult( - chunk_id=chunk.chunk_id, - doc_id=chunk.doc_id, - kb_id=chunk.kb_id, - content=chunk.content, + 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[:top_k] + # return results[: len(results) // len(kb_ids)] + return results[:top_k_sparse] diff --git a/astrbot/core/knowledge_base/session_config_db.py b/astrbot/core/knowledge_base/session_config_db.py deleted file mode 100644 index ce0c63a4..00000000 --- a/astrbot/core/knowledge_base/session_config_db.py +++ /dev/null @@ -1,157 +0,0 @@ -"""会话知识库配置数据库操作 - -该模块封装会话知识库配置的数据库查询操作。 - -注意: 会话配置表 (kb_session_config) 存储在知识库独立数据库 (kb.db) 中, - 而不是主数据库 (astrbot.db) 中,以实现完全解耦。 -""" - -import json -from typing import Optional - -from sqlalchemy import select - -from astrbot.core.knowledge_base.kb_sqlite import KBSQLiteDatabase -from astrbot.core.knowledge_base.models import KBSessionConfig - - -class SessionConfigDB: - """会话知识库配置数据库操作类 - - 职责: - - 提供会话知识库配置管理 - - 统一异常处理 - - 注意: 该类操作知识库独立数据库,实现完全解耦 - """ - - def __init__(self, db: KBSQLiteDatabase): - """初始化会话配置数据库操作类 - - Args: - db: 知识库独立数据库实例 (kb.db),不是主数据库 - """ - self.db = db - - async def get_session_kb_ids(self, session_id: str) -> list[str]: - """获取会话关联的知识库 ID 列表 - - 查找顺序: - 1. 会话级别配置 (优先) - 2. 平台级别配置 - 3. 返回空列表 - """ - async with self.db.get_db() as session: - # 1. 查找会话级别配置 - stmt = select(KBSessionConfig).where( - KBSessionConfig.scope == "session", - KBSessionConfig.scope_id == session_id, - ) - result = await session.execute(stmt) - config = result.scalar_one_or_none() - - if config: - return json.loads(config.kb_ids) - - # 2. 提取平台 ID (格式: platform:xxx:session_id) - parts = session_id.split(":") - if len(parts) >= 2: - platform_id = parts[0] - - # 查找平台级别配置 - stmt = select(KBSessionConfig).where( - KBSessionConfig.scope == "platform", - KBSessionConfig.scope_id == platform_id, - ) - result = await session.execute(stmt) - config = result.scalar_one_or_none() - - if config: - return json.loads(config.kb_ids) - - # 3. 无配置 - return [] - - async def set_session_kb_ids( - self, - scope: str, - scope_id: str, - kb_ids: list[str], - top_k: Optional[int] = None, - enable_rerank: Optional[bool] = None, - ) -> KBSessionConfig: - """设置会话知识库配置 - - Args: - scope: 配置范围 (session/platform) - scope_id: 范围标识 (会话 ID 或平台 ID) - kb_ids: 知识库 ID 列表 - top_k: 返回结果数量 (可选) - enable_rerank: 是否启用 Rerank (可选) - """ - async with self.db.get_db() as session: - # 查找现有配置 - stmt = select(KBSessionConfig).where( - KBSessionConfig.scope == scope, - KBSessionConfig.scope_id == scope_id, - ) - result = await session.execute(stmt) - config = result.scalar_one_or_none() - - if config: - # 更新现有配置 - config.kb_ids = json.dumps(kb_ids) - if top_k is not None: - config.top_k = top_k - if enable_rerank is not None: - config.enable_rerank = enable_rerank - else: - # 创建新配置 - config = KBSessionConfig( - scope=scope, - scope_id=scope_id, - kb_ids=json.dumps(kb_ids), - top_k=top_k, - enable_rerank=enable_rerank, - ) - session.add(config) - - await session.commit() - await session.refresh(config) - return config - - async def delete_session_kb_config(self, scope: str, scope_id: str) -> bool: - """删除会话知识库配置""" - async with self.db.get_db() as session: - stmt = select(KBSessionConfig).where( - KBSessionConfig.scope == scope, - KBSessionConfig.scope_id == scope_id, - ) - result = await session.execute(stmt) - config = result.scalar_one_or_none() - - if not config: - return False - - await session.delete(config) - await session.commit() - return True - - async def list_all_session_configs( - self, offset: int = 0, limit: int = 100, scope: Optional[str] = None - ) -> list[KBSessionConfig]: - """列出所有会话配置""" - async with self.db.get_db() as session: - stmt = select(KBSessionConfig) - - if scope: - stmt = stmt.where(KBSessionConfig.scope == scope) - - stmt = ( - stmt.offset(offset) - .limit(limit) - .order_by(KBSessionConfig.created_at.desc()) - ) - - result = await session.execute(stmt) - return list(result.scalars().all()) diff --git a/astrbot/core/pipeline/process_stage/method/llm_request.py b/astrbot/core/pipeline/process_stage/method/llm_request.py index c462e4db..703b3681 100644 --- a/astrbot/core/pipeline/process_stage/method/llm_request.py +++ b/astrbot/core/pipeline/process_stage/method/llm_request.py @@ -7,7 +7,7 @@ import copy import json import traceback from datetime import timedelta -from typing import AsyncGenerator, Union +from collections.abc import AsyncGenerator from astrbot.core.conversation_mgr import Conversation from astrbot.core import logger from astrbot.core.message.components import Image @@ -45,7 +45,7 @@ except (ModuleNotFoundError, ImportError): AgentContextWrapper = ContextWrapper[AstrAgentContext] -AgentRunner = ToolLoopAgentRunner[AgentContextWrapper] +AgentRunner = ToolLoopAgentRunner[AstrAgentContext] class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]): @@ -103,7 +103,7 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]): request = ProviderRequest( prompt=input_, - system_prompt=tool.description, + system_prompt=tool.description or "", image_urls=[], # 暂时不传递原始 agent 的上下文 contexts=[], # 暂时不传递原始 agent 的上下文 func_tool=toolset, @@ -240,7 +240,7 @@ class FunctionToolExecutor(BaseFunctionToolExecutor[AstrAgentContext]): yield res -class MainAgentHooks(BaseAgentRunHooks[AgentContextWrapper]): +class MainAgentHooks(BaseAgentRunHooks[AstrAgentContext]): async def on_agent_done(self, run_context, llm_response): # 执行事件钩子 await call_event_hook( @@ -338,7 +338,7 @@ class LLMRequestSubStage(Stage): self.conv_manager = ctx.plugin_manager.context.conversation_manager - def _select_provider(self, event: AstrMessageEvent) -> Provider | None: + def _select_provider(self, event: AstrMessageEvent): """选择使用的 LLM 提供商""" sel_provider = event.get_extra("selected_provider") _ctx = self.ctx.plugin_manager.context @@ -368,7 +368,7 @@ class LLMRequestSubStage(Stage): async def process( self, event: AstrMessageEvent, _nested: bool = False - ) -> Union[None, AsyncGenerator[None, None]]: + ) -> None | AsyncGenerator[None, None]: req: ProviderRequest | None = None if not self.ctx.astrbot_config["provider_settings"]["enable"]: @@ -383,6 +383,9 @@ class LLMRequestSubStage(Stage): provider = self._select_provider(event) if provider is None: return + if not isinstance(provider, Provider): + logger.error(f"选择的提供商类型无效({type(provider)})，跳过 LLM 请求处理。") + return if event.get_extra("provider_request"): req = event.get_extra("provider_request") @@ -489,6 +492,9 @@ class LLMRequestSubStage(Stage): new_tool_set.add_tool(tool) req.func_tool = new_tool_set + # 备份 req.contexts + backup_contexts = copy.deepcopy(req.contexts) + # run agent agent_runner = AgentRunner() logger.debug( @@ -526,8 +532,10 @@ class LLMRequestSubStage(Stage): chain = ( MessageChain().message(final_llm_resp.completion_text).chain ) - else: + elif final_llm_resp.result_chain: chain = final_llm_resp.result_chain.chain + else: + chain = MessageChain().chain event.set_result( MessageEventResult( chain=chain, @@ -538,6 +546,9 @@ class LLMRequestSubStage(Stage): async for _ in run_agent(agent_runner, self.max_step, self.show_tool_use): yield + # 恢复备份的 contexts + req.contexts = backup_contexts + await self._save_to_history(event, req, agent_runner.get_final_llm_resp()) # 异步处理 WebChat 特殊情况 @@ -556,6 +567,8 @@ class LLMRequestSubStage(Stage): self, event: AstrMessageEvent, req: ProviderRequest, prov: Provider ): """处理 WebChat 平台的特殊情况，包括第一次 LLM 对话时总结对话内容生成 title""" + if not req.conversation: + return conversation = await self.conv_manager.get_conversation( event.unified_msg_origin, req.conversation.cid ) diff --git a/astrbot/core/pipeline/process_stage/utils.py b/astrbot/core/pipeline/process_stage/utils.py index d870c398..7ba3757c 100644 --- a/astrbot/core/pipeline/process_stage/utils.py +++ b/astrbot/core/pipeline/process_stage/utils.py @@ -7,7 +7,6 @@ async def inject_kb_context( umo: str, p_ctx: PipelineContext, req: ProviderRequest, - top_k: int = 5, ) -> None: """inject knowledge base context into the provider request @@ -15,13 +14,19 @@ async def inject_kb_context( p_ctx: Pipeline context req: Provider request """ - kb_injector = p_ctx.plugin_manager.context.kb_manager.get_kb_injector() - if not kb_injector: + kb_mgr = p_ctx.plugin_manager.context.kb_manager + kb_names = p_ctx.astrbot_config.get("kb_names", []) + top_k_fusion = p_ctx.astrbot_config.get("kb_fusion_top_k", 20) + top_k = p_ctx.astrbot_config.get("kb_final_top_k", 5) + + if not kb_names: return - kb_context = await kb_injector.retrieve_and_inject( - unified_msg_origin=umo, + + kb_context = await kb_mgr.retrieve( query=req.prompt, - top_k=top_k, + kb_names=kb_names, + top_k_fusion=top_k_fusion, + top_m_final=top_k, ) if not kb_context: return diff --git a/astrbot/core/platform/astr_message_event.py b/astrbot/core/platform/astr_message_event.py index 48c63ef2..3a4b8c12 100644 --- a/astrbot/core/platform/astr_message_event.py +++ b/astrbot/core/platform/astr_message_event.py @@ -4,7 +4,7 @@ import re import hashlib import uuid -from typing import List, Union, Optional, AsyncGenerator, TypeVar, Any +from typing import List, Union, Optional, AsyncGenerator, Any from astrbot import logger from astrbot.core.db.po import Conversation @@ -26,8 +26,6 @@ from .astrbot_message import AstrBotMessage, Group from .platform_metadata import PlatformMetadata from .message_session import MessageSession, MessageSesion # noqa -_VT = TypeVar("_VT") - class AstrMessageEvent(abc.ABC): def __init__( @@ -92,8 +90,10 @@ class AstrMessageEvent(abc.ABC): """ return self.message_str - def _outline_chain(self, chain: List[BaseMessageComponent]) -> str: + def _outline_chain(self, chain: Optional[List[BaseMessageComponent]]) -> str: outline = "" + if not chain: + return outline for i in chain: if isinstance(i, Plain): outline += i.text @@ -175,9 +175,7 @@ class AstrMessageEvent(abc.ABC): """ self._extras[key] = value - def get_extra( - self, key: str | None = None, default: _VT = None - ) -> dict[str, Any] | _VT: + def get_extra(self, key: str | None = None, default=None) -> Any: """ 获取额外的信息。 """ @@ -265,6 +263,9 @@ class AstrMessageEvent(abc.ABC): """ if isinstance(result, str): result = MessageEventResult().message(result) + # 兼容外部插件或调用方传入的 chain=None 的情况，确保为可迭代列表 + if isinstance(result, MessageEventResult) and result.chain is None: + result.chain = [] self._result = result def stop_event(self): diff --git a/astrbot/core/platform/sources/misskey/misskey_adapter.py b/astrbot/core/platform/sources/misskey/misskey_adapter.py index 84608b54..981d05c8 100644 --- a/astrbot/core/platform/sources/misskey/misskey_adapter.py +++ b/astrbot/core/platform/sources/misskey/misskey_adapter.py @@ -1,6 +1,6 @@ import asyncio -import json -from typing import Dict, Any, Optional, Awaitable +import random +from typing import Dict, Any, Optional, Awaitable, List from astrbot.api import logger from astrbot.api.event import MessageChain @@ -14,6 +14,13 @@ from astrbot.core.platform.astr_message_event import MessageSession import astrbot.api.message_components as Comp from .misskey_api import MisskeyAPI +import os + +try: + import magic # type: ignore +except Exception: + magic = None + from .misskey_event import MisskeyPlatformEvent from .misskey_utils import ( serialize_message_chain, @@ -25,9 +32,15 @@ from .misskey_utils import ( extract_sender_info, create_base_message, process_at_mention, + format_poll, cache_user_info, cache_room_info, ) +from astrbot.core.utils.astrbot_path import get_astrbot_data_path + +# Constants +MAX_FILE_UPLOAD_COUNT = 16 +DEFAULT_UPLOAD_CONCURRENCY = 3 @register_platform_adapter("misskey", "Misskey 平台适配器") @@ -46,6 +59,31 @@ class MisskeyPlatformAdapter(Platform): ) self.local_only = self.config.get("misskey_local_only", False) self.enable_chat = self.config.get("misskey_enable_chat", True) + self.enable_file_upload = self.config.get("misskey_enable_file_upload", True) + self.upload_folder = self.config.get("misskey_upload_folder") + + # download / security related options (exposed to platform_config) + self.allow_insecure_downloads = bool( + self.config.get("misskey_allow_insecure_downloads", False) + ) + # parse download timeout and chunk size safely + _dt = self.config.get("misskey_download_timeout") + try: + self.download_timeout = int(_dt) if _dt is not None else 15 + except Exception: + self.download_timeout = 15 + + _chunk = self.config.get("misskey_download_chunk_size") + try: + self.download_chunk_size = int(_chunk) if _chunk is not None else 64 * 1024 + except Exception: + self.download_chunk_size = 64 * 1024 + # parse max download bytes safely + _md_bytes = self.config.get("misskey_max_download_bytes") + try: + self.max_download_bytes = int(_md_bytes) if _md_bytes is not None else None + except Exception: + self.max_download_bytes = None self.unique_session = platform_settings["unique_session"] @@ -63,6 +101,11 @@ class MisskeyPlatformAdapter(Platform): "misskey_default_visibility": "public", "misskey_local_only": False, "misskey_enable_chat": True, + # download / security options + "misskey_allow_insecure_downloads": False, + "misskey_download_timeout": 15, + "misskey_download_chunk_size": 65536, + "misskey_max_download_bytes": None, } default_config.update(self.config) @@ -78,7 +121,14 @@ class MisskeyPlatformAdapter(Platform): logger.error("[Misskey] 配置不完整，无法启动") return - self.api = MisskeyAPI(self.instance_url, self.access_token) + self.api = MisskeyAPI( + self.instance_url, + self.access_token, + allow_insecure_downloads=self.allow_insecure_downloads, + download_timeout=self.download_timeout, + chunk_size=self.download_chunk_size, + max_download_bytes=self.max_download_bytes, + ) self._running = True try: @@ -95,6 +145,80 @@ class MisskeyPlatformAdapter(Platform): await self._start_websocket_connection() + def _register_event_handlers(self, streaming): + """注册事件处理器""" + streaming.add_message_handler("notification", self._handle_notification) + streaming.add_message_handler("main:notification", self._handle_notification) + + if self.enable_chat: + streaming.add_message_handler("newChatMessage", self._handle_chat_message) + streaming.add_message_handler( + "messaging:newChatMessage", self._handle_chat_message + ) + streaming.add_message_handler("_debug", self._debug_handler) + + async def _send_text_only_message( + self, session_id: str, text: str, session, message_chain + ): + """发送纯文本消息（无文件上传）""" + if not self.api: + return await super().send_by_session(session, message_chain) + + if session_id and is_valid_user_session_id(session_id): + from .misskey_utils import extract_user_id_from_session_id + + user_id = extract_user_id_from_session_id(session_id) + payload: Dict[str, Any] = {"toUserId": user_id, "text": text} + await self.api.send_message(payload) + elif session_id and is_valid_room_session_id(session_id): + from .misskey_utils import extract_room_id_from_session_id + + room_id = extract_room_id_from_session_id(session_id) + payload = {"toRoomId": room_id, "text": text} + await self.api.send_room_message(payload) + + return await super().send_by_session(session, message_chain) + + def _process_poll_data( + self, message: AstrBotMessage, poll: Dict[str, Any], message_parts: List[str] + ): + """处理投票数据，将其添加到消息中""" + try: + if not isinstance(message.raw_message, dict): + message.raw_message = {} + message.raw_message["poll"] = poll + setattr(message, "poll", poll) + except Exception: + pass + + poll_text = format_poll(poll) + if poll_text: + message.message.append(Comp.Plain(poll_text)) + message_parts.append(poll_text) + + def _extract_additional_fields(self, session, message_chain) -> Dict[str, Any]: + """从会话和消息链中提取额外字段""" + fields = {"cw": None, "poll": None, "renote_id": None, "channel_id": None} + + for comp in message_chain.chain: + if hasattr(comp, "cw") and getattr(comp, "cw", None): + fields["cw"] = getattr(comp, "cw") + break + + if hasattr(session, "extra_data") and isinstance( + getattr(session, "extra_data", None), dict + ): + extra_data = getattr(session, "extra_data") + fields.update( + { + "poll": extra_data.get("poll"), + "renote_id": extra_data.get("renote_id"), + "channel_id": extra_data.get("channel_id"), + } + ) + + return fields + async def _start_websocket_connection(self): backoff_delay = 1.0 max_backoff = 300.0 @@ -109,25 +233,20 @@ class MisskeyPlatformAdapter(Platform): break streaming = self.api.get_streaming_client() - streaming.add_message_handler("notification", self._handle_notification) - if self.enable_chat: - streaming.add_message_handler( - "newChatMessage", self._handle_chat_message - ) - streaming.add_message_handler("_debug", self._debug_handler) + self._register_event_handlers(streaming) if await streaming.connect(): logger.info( f"[Misskey] WebSocket 已连接 (尝试 #{connection_attempts})" ) - connection_attempts = 0 # 重置计数器 + connection_attempts = 0 await streaming.subscribe_channel("main") if self.enable_chat: await streaming.subscribe_channel("messaging") await streaming.subscribe_channel("messagingIndex") logger.info("[Misskey] 聊天频道已订阅") - backoff_delay = 1.0 # 重置延迟 + backoff_delay = 1.0 await streaming.listen() else: logger.error( @@ -140,18 +259,20 @@ class MisskeyPlatformAdapter(Platform): ) if self._running: + jitter = random.uniform(0, 1.0) + sleep_time = backoff_delay + jitter logger.info( - f"[Misskey] {backoff_delay:.1f}秒后重连 (下次尝试 #{connection_attempts + 1})" + f"[Misskey] {sleep_time:.1f}秒后重连 (下次尝试 #{connection_attempts + 1})" ) - await asyncio.sleep(backoff_delay) + await asyncio.sleep(sleep_time) backoff_delay = min(backoff_delay * backoff_multiplier, max_backoff) async def _handle_notification(self, data: Dict[str, Any]): try: - logger.debug( - f"[Misskey] 收到通知事件:\n{json.dumps(data, indent=2, ensure_ascii=False)}" - ) notification_type = data.get("type") + logger.debug( + f"[Misskey] 收到通知事件: type={notification_type}, user_id={data.get('userId', 'unknown')}" + ) if notification_type in ["mention", "reply", "quote"]: note = data.get("note") if note and self._is_bot_mentioned(note): @@ -164,7 +285,7 @@ class MisskeyPlatformAdapter(Platform): message_obj=message, platform_meta=self.meta(), session_id=message.session_id, - client=self.api, + client=self, ) self.commit_event(event) except Exception as e: @@ -172,17 +293,16 @@ class MisskeyPlatformAdapter(Platform): async def _handle_chat_message(self, data: Dict[str, Any]): try: - logger.debug( - f"[Misskey] 收到聊天事件数据:\n{json.dumps(data, indent=2, ensure_ascii=False)}" - ) - sender_id = str( data.get("fromUserId", "") or data.get("fromUser", {}).get("id", "") ) + room_id = data.get("toRoomId") + logger.debug( + f"[Misskey] 收到聊天事件: sender_id={sender_id}, room_id={room_id}, is_self={sender_id == self.client_self_id}" + ) if sender_id == self.client_self_id: return - room_id = data.get("toRoomId") if room_id: raw_text = data.get("text", "") logger.debug( @@ -200,15 +320,16 @@ class MisskeyPlatformAdapter(Platform): message_obj=message, platform_meta=self.meta(), session_id=message.session_id, - client=self.api, + client=self, ) self.commit_event(event) except Exception as e: logger.error(f"[Misskey] 处理聊天消息失败: {e}") async def _debug_handler(self, data: Dict[str, Any]): + event_type = data.get("type", "unknown") logger.debug( - f"[Misskey] 收到未处理事件:\n{json.dumps(data, indent=2, ensure_ascii=False)}" + f"[Misskey] 收到未处理事件: type={event_type}, channel={data.get('channel', 'unknown')}" ) def _is_bot_mentioned(self, note: Dict[str, Any]) -> bool: @@ -239,43 +360,250 @@ class MisskeyPlatformAdapter(Platform): try: session_id = session.session_id + text, has_at_user = serialize_message_chain(message_chain.chain) if not has_at_user and session_id: - user_info = self._user_cache.get(session_id) + # 从session_id中提取用户ID用于缓存查询 + # session_id格式为: "chat%" 或 "room%" 或 "note%" + user_id_for_cache = None + if "%" in session_id: + parts = session_id.split("%") + if len(parts) >= 2: + user_id_for_cache = parts[1] + + user_info = None + if user_id_for_cache: + user_info = self._user_cache.get(user_id_for_cache) + text = add_at_mention_if_needed(text, user_info, has_at_user) + # 检查是否有文件组件 + has_file_components = any( + isinstance(comp, Comp.Image) + or isinstance(comp, Comp.File) + or hasattr(comp, "convert_to_file_path") + or hasattr(comp, "get_file") + or any( + hasattr(comp, a) for a in ("file", "url", "path", "src", "source") + ) + for comp in message_chain.chain + ) + if not text or not text.strip(): - logger.warning("[Misskey] 消息内容为空，跳过发送") - return await super().send_by_session(session, message_chain) + if not has_file_components: + logger.warning("[Misskey] 消息内容为空且无文件组件，跳过发送") + return await super().send_by_session(session, message_chain) + else: + text = "" if len(text) > self.max_message_length: text = text[: self.max_message_length] + "..." - if session_id and is_valid_user_session_id(session_id): - from .misskey_utils import extract_user_id_from_session_id + file_ids: List[str] = [] + fallback_urls: List[str] = [] - user_id = extract_user_id_from_session_id(session_id) - await self.api.send_message(user_id, text) - elif session_id and is_valid_room_session_id(session_id): + if not self.enable_file_upload: + return await self._send_text_only_message( + session_id, text, session, message_chain + ) + + MAX_UPLOAD_CONCURRENCY = 10 + upload_concurrency = int( + self.config.get( + "misskey_upload_concurrency", DEFAULT_UPLOAD_CONCURRENCY + ) + ) + upload_concurrency = min(upload_concurrency, MAX_UPLOAD_CONCURRENCY) + sem = asyncio.Semaphore(upload_concurrency) + + async def _upload_comp(comp) -> Optional[object]: + """组件上传函数：处理 URL（下载后上传）或本地文件（直接上传）""" + from .misskey_utils import ( + resolve_component_url_or_path, + upload_local_with_retries, + ) + + local_path = None + try: + async with sem: + if not self.api: + return None + + # 解析组件的 URL 或本地路径 + url_candidate, local_path = await resolve_component_url_or_path( + comp + ) + + if not url_candidate and not local_path: + return None + + preferred_name = getattr(comp, "name", None) or getattr( + comp, "file", None + ) + + # URL 上传：下载后本地上传 + if url_candidate: + result = await self.api.upload_and_find_file( + str(url_candidate), + preferred_name, + folder_id=self.upload_folder, + ) + if isinstance(result, dict) and result.get("id"): + return str(result["id"]) + + # 本地文件上传 + if local_path: + file_id = await upload_local_with_retries( + self.api, + str(local_path), + preferred_name, + self.upload_folder, + ) + if file_id: + return file_id + + # 所有上传都失败，尝试获取 URL 作为回退 + if hasattr(comp, "register_to_file_service"): + try: + url = await comp.register_to_file_service() + if url: + return {"fallback_url": url} + except Exception: + pass + + return None + + finally: + # 清理临时文件 + if local_path and isinstance(local_path, str): + data_temp = os.path.join(get_astrbot_data_path(), "temp") + if local_path.startswith(data_temp) and os.path.exists( + local_path + ): + try: + os.remove(local_path) + logger.debug(f"[Misskey] 已清理临时文件: {local_path}") + except Exception: + pass + + # 收集所有可能包含文件/URL信息的组件：支持异步接口或同步字段 + file_components = [] + for comp in message_chain.chain: + try: + if ( + isinstance(comp, Comp.Image) + or isinstance(comp, Comp.File) + or hasattr(comp, "convert_to_file_path") + or hasattr(comp, "get_file") + or any( + hasattr(comp, a) + for a in ("file", "url", "path", "src", "source") + ) + ): + file_components.append(comp) + except Exception: + # 保守跳过无法访问属性的组件 + continue + + if len(file_components) > MAX_FILE_UPLOAD_COUNT: + logger.warning( + f"[Misskey] 文件数量超过限制 ({len(file_components)} > {MAX_FILE_UPLOAD_COUNT})，只上传前{MAX_FILE_UPLOAD_COUNT}个文件" + ) + file_components = file_components[:MAX_FILE_UPLOAD_COUNT] + + upload_tasks = [_upload_comp(comp) for comp in file_components] + + try: + results = await asyncio.gather(*upload_tasks) if upload_tasks else [] + for r in results: + if not r: + continue + if isinstance(r, dict) and r.get("fallback_url"): + url = r.get("fallback_url") + if url: + fallback_urls.append(str(url)) + else: + try: + fid_str = str(r) + if fid_str: + file_ids.append(fid_str) + except Exception: + pass + except Exception: + logger.debug("[Misskey] 并发上传过程中出现异常，继续发送文本") + + if session_id and is_valid_room_session_id(session_id): from .misskey_utils import extract_room_id_from_session_id room_id = extract_room_id_from_session_id(session_id) - await self.api.send_room_message(room_id, text) - else: - visibility, visible_user_ids = resolve_message_visibility( - user_id=session_id, - user_cache=self._user_cache, - self_id=self.client_self_id, - default_visibility=self.default_visibility, + if fallback_urls: + appended = "\n" + "\n".join(fallback_urls) + text = (text or "") + appended + payload: Dict[str, Any] = {"toRoomId": room_id, "text": text} + if file_ids: + payload["fileIds"] = file_ids + await self.api.send_room_message(payload) + elif session_id: + from .misskey_utils import ( + extract_user_id_from_session_id, + is_valid_chat_session_id, ) - await self.api.create_note( - text, - visibility=visibility, - visible_user_ids=visible_user_ids, - local_only=self.local_only, - ) + if is_valid_chat_session_id(session_id): + user_id = extract_user_id_from_session_id(session_id) + if fallback_urls: + appended = "\n" + "\n".join(fallback_urls) + text = (text or "") + appended + payload: Dict[str, Any] = {"toUserId": user_id, "text": text} + if file_ids: + # 聊天消息只支持单个文件，使用 fileId 而不是 fileIds + payload["fileId"] = file_ids[0] + if len(file_ids) > 1: + logger.warning( + f"[Misskey] 聊天消息只支持单个文件，忽略其余 {len(file_ids) - 1} 个文件" + ) + await self.api.send_message(payload) + else: + # 回退到发帖逻辑 + # 去掉 session_id 中的 note% 前缀以匹配 user_cache 的键格式 + user_id_for_cache = ( + session_id.split("%")[1] if "%" in session_id else session_id + ) + + # 获取用户缓存信息（包含reply_to_note_id） + user_info_for_reply = self._user_cache.get(user_id_for_cache, {}) + + visibility, visible_user_ids = resolve_message_visibility( + user_id=user_id_for_cache, + user_cache=self._user_cache, + self_id=self.client_self_id, + default_visibility=self.default_visibility, + ) + logger.debug( + f"[Misskey] 解析可见性: visibility={visibility}, visible_user_ids={visible_user_ids}, session_id={session_id}, user_id_for_cache={user_id_for_cache}" + ) + + fields = self._extract_additional_fields(session, message_chain) + if fallback_urls: + appended = "\n" + "\n".join(fallback_urls) + text = (text or "") + appended + + # 从缓存中获取原消息ID作为reply_id + reply_id = user_info_for_reply.get("reply_to_note_id") + + await self.api.create_note( + text=text, + visibility=visibility, + visible_user_ids=visible_user_ids, + file_ids=file_ids or None, + local_only=self.local_only, + reply_id=reply_id, # 添加reply_id参数 + cw=fields["cw"], + poll=fields["poll"], + renote_id=fields["renote_id"], + channel_id=fields["channel_id"], + ) except Exception as e: logger.error(f"[Misskey] 发送消息失败: {e}") @@ -309,6 +637,14 @@ class MisskeyPlatformAdapter(Platform): file_parts = process_files(message, files) message_parts.extend(file_parts) + poll = raw_data.get("poll") or ( + raw_data.get("note", {}).get("poll") + if isinstance(raw_data.get("note"), dict) + else None + ) + if poll and isinstance(poll, dict): + self._process_poll_data(message, poll, message_parts) + message.message_str = ( " ".join(part for part in message_parts if part.strip()) if message_parts diff --git a/astrbot/core/platform/sources/misskey/misskey_api.py b/astrbot/core/platform/sources/misskey/misskey_api.py index dc4adcdd..4b920508 100644 --- a/astrbot/core/platform/sources/misskey/misskey_api.py +++ b/astrbot/core/platform/sources/misskey/misskey_api.py @@ -1,4 +1,6 @@ import json +import random +import asyncio from typing import Any, Optional, Dict, List, Callable, Awaitable import uuid @@ -11,6 +13,7 @@ except ImportError as e: ) from e from astrbot.api import logger +from .misskey_utils import FileIDExtractor # Constants API_MAX_RETRIES = 3 @@ -55,6 +58,7 @@ class StreamingClient: self.is_connected = False self.message_handlers: Dict[str, Callable] = {} self.channels: Dict[str, str] = {} + self.desired_channels: Dict[str, Optional[Dict]] = {} self._running = False self._last_pong = None @@ -72,6 +76,18 @@ class StreamingClient: self._running = True logger.info("[Misskey WebSocket] 已连接") + if self.desired_channels: + try: + desired = list(self.desired_channels.items()) + for channel_type, params in desired: + try: + await self.subscribe_channel(channel_type, params) + except Exception as e: + logger.warning( + f"[Misskey WebSocket] 重新订阅 {channel_type} 失败: {e}" + ) + except Exception: + pass return True except Exception as e: @@ -112,9 +128,12 @@ class StreamingClient: return message = {"type": "disconnect", "body": {"id": channel_id}} - await self.websocket.send(json.dumps(message)) - del self.channels[channel_id] + channel_type = self.channels.get(channel_id) + if channel_id in self.channels: + del self.channels[channel_id] + if channel_type and channel_type not in self.channels.values(): + self.desired_channels.pop(channel_type, None) def add_message_handler( self, event_type: str, handler: Callable[[Dict], Awaitable[None]] @@ -141,25 +160,67 @@ class StreamingClient: except websockets.exceptions.ConnectionClosedError as e: logger.warning(f"[Misskey WebSocket] 连接意外关闭: {e}") self.is_connected = False + try: + await self.disconnect() + except Exception: + pass except websockets.exceptions.ConnectionClosed as e: logger.warning( f"[Misskey WebSocket] 连接已关闭 (代码: {e.code}, 原因: {e.reason})" ) self.is_connected = False + try: + await self.disconnect() + except Exception: + pass except websockets.exceptions.InvalidHandshake as e: logger.error(f"[Misskey WebSocket] 握手失败: {e}") self.is_connected = False + try: + await self.disconnect() + except Exception: + pass except Exception as e: logger.error(f"[Misskey WebSocket] 监听消息失败: {e}") self.is_connected = False + try: + await self.disconnect() + except Exception: + pass async def _handle_message(self, data: Dict[str, Any]): message_type = data.get("type") body = data.get("body", {}) - logger.debug( - f"[Misskey WebSocket] 收到消息类型: {message_type}\n数据: {json.dumps(data, indent=2, ensure_ascii=False)}" - ) + def _build_channel_summary(message_type: Optional[str], body: Any) -> str: + try: + if not isinstance(body, dict): + return f"[Misskey WebSocket] 收到消息类型: {message_type}" + + inner = body.get("body") if isinstance(body.get("body"), dict) else body + note = ( + inner.get("note") + if isinstance(inner, dict) and isinstance(inner.get("note"), dict) + else None + ) + + text = note.get("text") if note else None + note_id = note.get("id") if note else None + files = note.get("files") or [] if note else [] + has_files = bool(files) + is_hidden = bool(note.get("isHidden")) if note else False + user = note.get("user", {}) if note else None + + return ( + f"[Misskey WebSocket] 收到消息类型: {message_type} | " + f"note_id={note_id} | user={user.get('username') if user else None} | " + f"text={text[:80] if text else '[no-text]'} | files={has_files} | hidden={is_hidden}" + ) + except Exception: + return f"[Misskey WebSocket] 收到消息类型: {message_type}" + + channel_summary = _build_channel_summary(message_type, body) + logger.info(channel_summary) if message_type == "channel": channel_id = body.get("id") @@ -202,16 +263,60 @@ class StreamingClient: await self.message_handlers["_debug"](data) -def retry_async(max_retries: int = 3, retryable_exceptions: tuple = ()): +def retry_async( + max_retries: int = 3, + retryable_exceptions: tuple = (APIConnectionError, APIRateLimitError), + backoff_base: float = 1.0, + max_backoff: float = 30.0, +): + """ + 智能异步重试装饰器 + + Args: + max_retries: 最大重试次数 + retryable_exceptions: 可重试的异常类型 + backoff_base: 退避基数 + max_backoff: 最大退避时间 + """ + def decorator(func): async def wrapper(*args, **kwargs): last_exc = None - for _ in range(max_retries): + func_name = getattr(func, "__name__", "unknown") + + for attempt in range(1, max_retries + 1): try: return await func(*args, **kwargs) except retryable_exceptions as e: last_exc = e + if attempt == max_retries: + logger.error( + f"[Misskey API] {func_name} 重试 {max_retries} 次后仍失败: {e}" + ) + break + + # 智能退避策略 + if isinstance(e, APIRateLimitError): + # 频率限制用更长的退避时间 + backoff = min(backoff_base * (3**attempt), max_backoff) + else: + # 其他错误用指数退避 + backoff = min(backoff_base * (2**attempt), max_backoff) + + jitter = random.uniform(0.1, 0.5) # 随机抖动 + sleep_time = backoff + jitter + + logger.warning( + f"[Misskey API] {func_name} 第 {attempt} 次重试失败: {e}，" + f"{sleep_time:.1f}s后重试" + ) + await asyncio.sleep(sleep_time) continue + except Exception as e: + # 非可重试异常直接抛出 + logger.error(f"[Misskey API] {func_name} 遇到不可重试异常: {e}") + raise + if last_exc: raise last_exc @@ -221,11 +326,27 @@ def retry_async(max_retries: int = 3, retryable_exceptions: tuple = ()): class MisskeyAPI: - def __init__(self, instance_url: str, access_token: str): + def __init__( + self, + instance_url: str, + access_token: str, + *, + allow_insecure_downloads: bool = False, + download_timeout: int = 15, + chunk_size: int = 64 * 1024, + max_download_bytes: Optional[int] = None, + ): self.instance_url = instance_url.rstrip("/") self.access_token = access_token self._session: Optional[aiohttp.ClientSession] = None self.streaming: Optional[StreamingClient] = None + # download options + self.allow_insecure_downloads = allow_insecure_downloads + self.download_timeout = download_timeout + self.chunk_size = chunk_size + self.max_download_bytes = ( + int(max_download_bytes) if max_download_bytes is not None else None + ) async def __aenter__(self): return self @@ -258,16 +379,37 @@ class MisskeyAPI: def _handle_response_status(self, status: int, endpoint: str): """处理 HTTP 响应状态码""" if status == 400: - logger.error(f"API 请求错误: {endpoint} (状态码: {status})") + logger.error(f"[Misskey API] 请求参数错误: {endpoint} (HTTP {status})") raise APIError(f"Bad request for {endpoint}") - elif status in (401, 403): - logger.error(f"API 认证失败: {endpoint} (状态码: {status})") - raise AuthenticationError(f"Authentication failed for {endpoint}") + elif status == 401: + logger.error(f"[Misskey API] 未授权访问: {endpoint} (HTTP {status})") + raise AuthenticationError(f"Unauthorized access for {endpoint}") + elif status == 403: + logger.error(f"[Misskey API] 访问被禁止: {endpoint} (HTTP {status})") + raise AuthenticationError(f"Forbidden access for {endpoint}") + elif status == 404: + logger.error(f"[Misskey API] 资源不存在: {endpoint} (HTTP {status})") + raise APIError(f"Resource not found for {endpoint}") + elif status == 413: + logger.error(f"[Misskey API] 请求体过大: {endpoint} (HTTP {status})") + raise APIError(f"Request entity too large for {endpoint}") elif status == 429: - logger.warning(f"API 频率限制: {endpoint} (状态码: {status})") + logger.warning(f"[Misskey API] 请求频率限制: {endpoint} (HTTP {status})") raise APIRateLimitError(f"Rate limit exceeded for {endpoint}") + elif status == 500: + logger.error(f"[Misskey API] 服务器内部错误: {endpoint} (HTTP {status})") + raise APIConnectionError(f"Internal server error for {endpoint}") + elif status == 502: + logger.error(f"[Misskey API] 网关错误: {endpoint} (HTTP {status})") + raise APIConnectionError(f"Bad gateway for {endpoint}") + elif status == 503: + logger.error(f"[Misskey API] 服务不可用: {endpoint} (HTTP {status})") + raise APIConnectionError(f"Service unavailable for {endpoint}") + elif status == 504: + logger.error(f"[Misskey API] 网关超时: {endpoint} (HTTP {status})") + raise APIConnectionError(f"Gateway timeout for {endpoint}") else: - logger.error(f"API 请求失败: {endpoint} (状态码: {status})") + logger.error(f"[Misskey API] 未知错误: {endpoint} (HTTP {status})") raise APIConnectionError(f"HTTP {status} for {endpoint}") async def _process_response( @@ -286,21 +428,25 @@ class MisskeyAPI: else [] ) if notifications_data: - logger.debug(f"获取到 {len(notifications_data)} 条新通知") + logger.debug( + f"[Misskey API] 获取到 {len(notifications_data)} 条新通知" + ) else: - logger.debug(f"API 请求成功: {endpoint}") + logger.debug(f"[Misskey API] 请求成功: {endpoint}") return result except json.JSONDecodeError as e: - logger.error(f"响应不是有效的 JSON 格式: {e}") + logger.error(f"[Misskey API] 响应格式错误: {e}") raise APIConnectionError("Invalid JSON response") from e else: try: error_text = await response.text() logger.error( - f"API 请求失败: {endpoint} - 状态码: {response.status}, 响应: {error_text}" + f"[Misskey API] 请求失败: {endpoint} - HTTP {response.status}, 响应: {error_text}" ) except Exception: - logger.error(f"API 请求失败: {endpoint} - 状态码: {response.status}") + logger.error( + f"[Misskey API] 请求失败: {endpoint} - HTTP {response.status}" + ) self._handle_response_status(response.status, endpoint) raise APIConnectionError(f"Request failed for {endpoint}") @@ -321,53 +467,307 @@ class MisskeyAPI: async with self.session.post(url, json=payload) as response: return await self._process_response(response, endpoint) except aiohttp.ClientError as e: - logger.error(f"HTTP 请求错误: {e}") + logger.error(f"[Misskey API] HTTP 请求错误: {e}") raise APIConnectionError(f"HTTP request failed: {e}") from e async def create_note( self, - text: str, + text: Optional[str] = None, visibility: str = "public", reply_id: Optional[str] = None, visible_user_ids: Optional[List[str]] = None, + file_ids: Optional[List[str]] = None, local_only: bool = False, + cw: Optional[str] = None, + poll: Optional[Dict[str, Any]] = None, + renote_id: Optional[str] = None, + channel_id: Optional[str] = None, + reaction_acceptance: Optional[str] = None, + no_extract_mentions: Optional[bool] = None, + no_extract_hashtags: Optional[bool] = None, + no_extract_emojis: Optional[bool] = None, + media_ids: Optional[List[str]] = None, ) -> Dict[str, Any]: - """创建新贴文""" - data: Dict[str, Any] = { - "text": text, - "visibility": visibility, - "localOnly": local_only, - } + """Create a note (wrapper for notes/create). All additional fields are optional and passed through to the API.""" + data: Dict[str, Any] = {} + + if text is not None: + data["text"] = text + + data["visibility"] = visibility + data["localOnly"] = local_only + if reply_id: data["replyId"] = reply_id + if visible_user_ids and visibility == "specified": data["visibleUserIds"] = visible_user_ids + if file_ids: + data["fileIds"] = file_ids + if media_ids: + data["mediaIds"] = media_ids + + if cw is not None: + data["cw"] = cw + if poll is not None: + data["poll"] = poll + if renote_id is not None: + data["renoteId"] = renote_id + if channel_id is not None: + data["channelId"] = channel_id + if reaction_acceptance is not None: + data["reactionAcceptance"] = reaction_acceptance + if no_extract_mentions is not None: + data["noExtractMentions"] = bool(no_extract_mentions) + if no_extract_hashtags is not None: + data["noExtractHashtags"] = bool(no_extract_hashtags) + if no_extract_emojis is not None: + data["noExtractEmojis"] = bool(no_extract_emojis) + result = await self._make_request("notes/create", data) - note_id = result.get("createdNote", {}).get("id", "unknown") - logger.debug(f"发帖成功，note_id: {note_id}") + note_id = ( + result.get("createdNote", {}).get("id", "unknown") + if isinstance(result, dict) + else "unknown" + ) + logger.debug(f"[Misskey API] 发帖成功: {note_id}") return result + async def upload_file( + self, + file_path: str, + name: Optional[str] = None, + folder_id: Optional[str] = None, + ) -> Dict[str, Any]: + """Upload a file to Misskey drive/files/create and return a dict containing id and raw result.""" + if not file_path: + raise APIError("No file path provided for upload") + + url = f"{self.instance_url}/api/drive/files/create" + form = aiohttp.FormData() + form.add_field("i", self.access_token) + + try: + filename = name or file_path.split("/")[-1] + if folder_id: + form.add_field("folderId", str(folder_id)) + + try: + f = open(file_path, "rb") + except FileNotFoundError as e: + logger.error(f"[Misskey API] 本地文件不存在: {file_path}") + raise APIError(f"File not found: {file_path}") from e + + try: + form.add_field("file", f, filename=filename) + async with self.session.post(url, data=form) as resp: + result = await self._process_response(resp, "drive/files/create") + file_id = FileIDExtractor.extract_file_id(result) + logger.debug( + f"[Misskey API] 本地文件上传成功: {filename} -> {file_id}" + ) + return {"id": file_id, "raw": result} + finally: + f.close() + except aiohttp.ClientError as e: + logger.error(f"[Misskey API] 文件上传网络错误: {e}") + raise APIConnectionError(f"Upload failed: {e}") from e + + async def find_files_by_hash(self, md5_hash: str) -> List[Dict[str, Any]]: + """Find files by MD5 hash""" + if not md5_hash: + raise APIError("No MD5 hash provided for find-by-hash") + + data = {"md5": md5_hash} + + try: + logger.debug(f"[Misskey API] find-by-hash 请求: md5={md5_hash}") + result = await self._make_request("drive/files/find-by-hash", data) + logger.debug( + f"[Misskey API] find-by-hash 响应: 找到 {len(result) if isinstance(result, list) else 0} 个文件" + ) + return result if isinstance(result, list) else [] + except Exception as e: + logger.error(f"[Misskey API] 根据哈希查找文件失败: {e}") + raise + + async def find_files_by_name( + self, name: str, folder_id: Optional[str] = None + ) -> List[Dict[str, Any]]: + """Find files by name""" + if not name: + raise APIError("No name provided for find") + + data: Dict[str, Any] = {"name": name} + if folder_id: + data["folderId"] = folder_id + + try: + logger.debug(f"[Misskey API] find 请求: name={name}, folder_id={folder_id}") + result = await self._make_request("drive/files/find", data) + logger.debug( + f"[Misskey API] find 响应: 找到 {len(result) if isinstance(result, list) else 0} 个文件" + ) + return result if isinstance(result, list) else [] + except Exception as e: + logger.error(f"[Misskey API] 根据名称查找文件失败: {e}") + raise + + async def find_files( + self, + limit: int = 10, + folder_id: Optional[str] = None, + type: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """List files with optional filters""" + data: Dict[str, Any] = {"limit": limit} + if folder_id is not None: + data["folderId"] = folder_id + if type is not None: + data["type"] = type + + try: + logger.debug( + f"[Misskey API] 列表文件请求: limit={limit}, folder_id={folder_id}, type={type}" + ) + result = await self._make_request("drive/files", data) + logger.debug( + f"[Misskey API] 列表文件响应: 找到 {len(result) if isinstance(result, list) else 0} 个文件" + ) + return result if isinstance(result, list) else [] + except Exception as e: + logger.error(f"[Misskey API] 列表文件失败: {e}") + raise + + async def _download_with_existing_session( + self, url: str, ssl_verify: bool = True + ) -> Optional[bytes]: + """使用现有会话下载文件""" + if not (hasattr(self, "session") and self.session): + raise APIConnectionError("No existing session available") + + async with self.session.get( + url, timeout=aiohttp.ClientTimeout(total=15), ssl=ssl_verify + ) as response: + if response.status == 200: + return await response.read() + return None + + async def _download_with_temp_session( + self, url: str, ssl_verify: bool = True + ) -> Optional[bytes]: + """使用临时会话下载文件""" + connector = aiohttp.TCPConnector(ssl=ssl_verify) + async with aiohttp.ClientSession(connector=connector) as temp_session: + async with temp_session.get( + url, timeout=aiohttp.ClientTimeout(total=15) + ) as response: + if response.status == 200: + return await response.read() + return None + + async def upload_and_find_file( + self, + url: str, + name: Optional[str] = None, + folder_id: Optional[str] = None, + max_wait_time: float = 30.0, + check_interval: float = 2.0, + ) -> Optional[Dict[str, Any]]: + """ + 简化的文件上传：尝试 URL 上传，失败则下载后本地上传 + + Args: + url: 文件URL + name: 文件名（可选） + folder_id: 文件夹ID（可选） + max_wait_time: 保留参数（未使用） + check_interval: 保留参数（未使用） + + Returns: + 包含文件ID和元信息的字典，失败时返回None + """ + if not url: + raise APIError("URL不能为空") + + # 通过本地上传获取即时文件 ID（下载文件 → 上传 → 返回 ID） + try: + import tempfile + import os + + # SSL 验证下载，失败则重试不验证 SSL + tmp_bytes = None + try: + tmp_bytes = await self._download_with_existing_session( + url, ssl_verify=True + ) or await self._download_with_temp_session(url, ssl_verify=True) + except Exception as ssl_error: + logger.debug( + f"[Misskey API] SSL 验证下载失败: {ssl_error}，重试不验证 SSL" + ) + try: + tmp_bytes = await self._download_with_existing_session( + url, ssl_verify=False + ) or await self._download_with_temp_session(url, ssl_verify=False) + except Exception: + pass + + if tmp_bytes: + with tempfile.NamedTemporaryFile(delete=False) as tmpf: + tmpf.write(tmp_bytes) + tmp_path = tmpf.name + + try: + result = await self.upload_file(tmp_path, name, folder_id) + logger.debug(f"[Misskey API] 本地上传成功: {result.get('id')}") + return result + finally: + try: + os.unlink(tmp_path) + except Exception: + pass + except Exception as e: + logger.error(f"[Misskey API] 本地上传失败: {e}") + + return None + async def get_current_user(self) -> Dict[str, Any]: """获取当前用户信息""" return await self._make_request("i", {}) - async def send_message(self, user_id: str, text: str) -> Dict[str, Any]: - """发送聊天消息""" - result = await self._make_request( - "chat/messages/create-to-user", {"toUserId": user_id, "text": text} - ) + async def send_message( + self, user_id_or_payload: Any, text: Optional[str] = None + ) -> Dict[str, Any]: + """发送聊天消息。 + + Accepts either (user_id: str, text: str) or a single dict payload prepared by caller. + """ + if isinstance(user_id_or_payload, dict): + data = user_id_or_payload + else: + data = {"toUserId": user_id_or_payload, "text": text} + + result = await self._make_request("chat/messages/create-to-user", data) message_id = result.get("id", "unknown") - logger.debug(f"聊天发送成功，message_id: {message_id}") + logger.debug(f"[Misskey API] 聊天消息发送成功: {message_id}") return result - async def send_room_message(self, room_id: str, text: str) -> Dict[str, Any]: - """发送房间消息""" - result = await self._make_request( - "chat/messages/create-to-room", {"toRoomId": room_id, "text": text} - ) + async def send_room_message( + self, room_id_or_payload: Any, text: Optional[str] = None + ) -> Dict[str, Any]: + """发送房间消息。 + + Accepts either (room_id: str, text: str) or a single dict payload. + """ + if isinstance(room_id_or_payload, dict): + data = room_id_or_payload + else: + data = {"toRoomId": room_id_or_payload, "text": text} + + result = await self._make_request("chat/messages/create-to-room", data) message_id = result.get("id", "unknown") - logger.debug(f"房间消息发送成功，message_id: {message_id}") + logger.debug(f"[Misskey API] 房间消息发送成功: {message_id}") return result async def get_messages( @@ -381,9 +781,8 @@ class MisskeyAPI: result = await self._make_request("chat/messages/user-timeline", data) if isinstance(result, list): return result - else: - logger.warning(f"获取聊天消息响应格式异常: {type(result)}") - return [] + logger.warning(f"[Misskey API] 聊天消息响应格式异常: {type(result)}") + return [] async def get_mentions( self, limit: int = 10, since_id: Optional[str] = None @@ -400,5 +799,142 @@ class MisskeyAPI: elif isinstance(result, dict) and "notifications" in result: return result["notifications"] else: - logger.warning(f"获取提及通知响应格式异常: {type(result)}") + logger.warning(f"[Misskey API] 提及通知响应格式异常: {type(result)}") return [] + + async def send_message_with_media( + self, + message_type: str, + target_id: str, + text: Optional[str] = None, + media_urls: Optional[List[str]] = None, + local_files: Optional[List[str]] = None, + **kwargs, + ) -> Dict[str, Any]: + """ + 通用消息发送函数：统一处理文本+媒体发送 + + Args: + message_type: 消息类型 ('chat', 'room', 'note') + target_id: 目标ID (用户ID/房间ID/频道ID等) + text: 文本内容 + media_urls: 媒体文件URL列表 + local_files: 本地文件路径列表 + **kwargs: 其他参数（如visibility等） + + Returns: + 发送结果字典 + + Raises: + APIError: 参数错误或发送失败 + """ + if not text and not media_urls and not local_files: + raise APIError("消息内容不能为空：需要文本或媒体文件") + + file_ids = [] + + # 处理远程媒体文件 + if media_urls: + file_ids.extend(await self._process_media_urls(media_urls)) + + # 处理本地文件 + if local_files: + file_ids.extend(await self._process_local_files(local_files)) + + # 根据消息类型发送 + return await self._dispatch_message( + message_type, target_id, text, file_ids, **kwargs + ) + + async def _process_media_urls(self, urls: List[str]) -> List[str]: + """处理远程媒体文件URL列表，返回文件ID列表""" + file_ids = [] + for url in urls: + try: + result = await self.upload_and_find_file(url) + if result and result.get("id"): + file_ids.append(result["id"]) + logger.debug(f"[Misskey API] URL媒体上传成功: {result['id']}") + else: + logger.error(f"[Misskey API] URL媒体上传失败: {url}") + except Exception as e: + logger.error(f"[Misskey API] URL媒体处理失败 {url}: {e}") + # 继续处理其他文件，不中断整个流程 + continue + return file_ids + + async def _process_local_files(self, file_paths: List[str]) -> List[str]: + """处理本地文件路径列表，返回文件ID列表""" + file_ids = [] + for file_path in file_paths: + try: + result = await self.upload_file(file_path) + if result and result.get("id"): + file_ids.append(result["id"]) + logger.debug(f"[Misskey API] 本地文件上传成功: {result['id']}") + else: + logger.error(f"[Misskey API] 本地文件上传失败: {file_path}") + except Exception as e: + logger.error(f"[Misskey API] 本地文件处理失败 {file_path}: {e}") + continue + return file_ids + + async def _dispatch_message( + self, + message_type: str, + target_id: str, + text: Optional[str], + file_ids: List[str], + **kwargs, + ) -> Dict[str, Any]: + """根据消息类型分发到对应的发送方法""" + if message_type == "chat": + # 聊天消息使用 fileId (单数) + payload = {"toUserId": target_id} + if text: + payload["text"] = text + if file_ids: + if len(file_ids) == 1: + payload["fileId"] = file_ids[0] + else: + # 多文件时逐个发送 + results = [] + for file_id in file_ids: + single_payload = payload.copy() + single_payload["fileId"] = file_id + result = await self.send_message(single_payload) + results.append(result) + return {"multiple": True, "results": results} + return await self.send_message(payload) + + elif message_type == "room": + # 房间消息使用 fileId (单数) + payload = {"toRoomId": target_id} + if text: + payload["text"] = text + if file_ids: + if len(file_ids) == 1: + payload["fileId"] = file_ids[0] + else: + # 多文件时逐个发送 + results = [] + for file_id in file_ids: + single_payload = payload.copy() + single_payload["fileId"] = file_id + result = await self.send_room_message(single_payload) + results.append(result) + return {"multiple": True, "results": results} + return await self.send_room_message(payload) + + elif message_type == "note": + # 发帖使用 fileIds (复数) + note_kwargs = { + "text": text, + "file_ids": file_ids or None, + } + # 合并其他参数 + note_kwargs.update(kwargs) + return await self.create_note(**note_kwargs) + + else: + raise APIError(f"不支持的消息类型: {message_type}") diff --git a/astrbot/core/platform/sources/misskey/misskey_event.py b/astrbot/core/platform/sources/misskey/misskey_event.py index 391d10b5..cd737f78 100644 --- a/astrbot/core/platform/sources/misskey/misskey_event.py +++ b/astrbot/core/platform/sources/misskey/misskey_event.py @@ -40,48 +40,83 @@ class MisskeyPlatformEvent(AstrMessageEvent): return any(message_trimmed.startswith(prefix) for prefix in system_prefixes) async def send(self, message: MessageChain): - content, has_at = serialize_message_chain(message.chain) - - if not content: - logger.debug("[MisskeyEvent] 内容为空，跳过发送") - return - + """发送消息，使用适配器的完整上传和发送逻辑""" try: - original_message_id = getattr(self.message_obj, "message_id", None) - raw_message = getattr(self.message_obj, "raw_message", {}) + logger.debug( + f"[MisskeyEvent] send 方法被调用，消息链包含 {len(message.chain)} 个组件" + ) - if raw_message and not has_at: - user_data = raw_message.get("user", {}) - user_info = { - "username": user_data.get("username", ""), - "nickname": user_data.get("name", user_data.get("username", "")), - } - content = add_at_mention_if_needed(content, user_info, has_at) + # 使用适配器的 send_by_session 方法，它包含文件上传逻辑 + from astrbot.core.platform.message_session import MessageSession + from astrbot.core.platform.message_type import MessageType - # 根据会话类型选择发送方式 - if hasattr(self.client, "send_message") and is_valid_user_session_id( - self.session_id - ): - user_id = extract_user_id_from_session_id(self.session_id) - await self.client.send_message(user_id, content) - elif hasattr(self.client, "send_room_message") and is_valid_room_session_id( - self.session_id - ): - room_id = extract_room_id_from_session_id(self.session_id) - await self.client.send_room_message(room_id, content) - elif original_message_id and hasattr(self.client, "create_note"): - visibility, visible_user_ids = resolve_visibility_from_raw_message( - raw_message - ) - await self.client.create_note( - content, - reply_id=original_message_id, - visibility=visibility, - visible_user_ids=visible_user_ids, - ) - elif hasattr(self.client, "create_note"): - logger.debug("[MisskeyEvent] 创建新帖子") - await self.client.create_note(content) + # 根据session_id类型确定消息类型 + if is_valid_user_session_id(self.session_id): + message_type = MessageType.FRIEND_MESSAGE + elif is_valid_room_session_id(self.session_id): + message_type = MessageType.GROUP_MESSAGE + else: + message_type = MessageType.FRIEND_MESSAGE # 默认 + + session = MessageSession( + platform_name=self.platform_meta.name, + message_type=message_type, + session_id=self.session_id, + ) + + logger.debug( + f"[MisskeyEvent] 检查适配器方法: hasattr(self.client, 'send_by_session') = {hasattr(self.client, 'send_by_session')}" + ) + + # 调用适配器的 send_by_session 方法 + if hasattr(self.client, "send_by_session"): + logger.debug("[MisskeyEvent] 调用适配器的 send_by_session 方法") + await self.client.send_by_session(session, message) + else: + # 回退到原来的简化发送逻辑 + content, has_at = serialize_message_chain(message.chain) + + if not content: + logger.debug("[MisskeyEvent] 内容为空，跳过发送") + return + + original_message_id = getattr(self.message_obj, "message_id", None) + raw_message = getattr(self.message_obj, "raw_message", {}) + + if raw_message and not has_at: + user_data = raw_message.get("user", {}) + user_info = { + "username": user_data.get("username", ""), + "nickname": user_data.get( + "name", user_data.get("username", "") + ), + } + content = add_at_mention_if_needed(content, user_info, has_at) + + # 根据会话类型选择发送方式 + if hasattr(self.client, "send_message") and is_valid_user_session_id( + self.session_id + ): + user_id = extract_user_id_from_session_id(self.session_id) + await self.client.send_message(user_id, content) + elif hasattr( + self.client, "send_room_message" + ) and is_valid_room_session_id(self.session_id): + room_id = extract_room_id_from_session_id(self.session_id) + await self.client.send_room_message(room_id, content) + elif original_message_id and hasattr(self.client, "create_note"): + visibility, visible_user_ids = resolve_visibility_from_raw_message( + raw_message + ) + await self.client.create_note( + content, + reply_id=original_message_id, + visibility=visibility, + visible_user_ids=visible_user_ids, + ) + elif hasattr(self.client, "create_note"): + logger.debug("[MisskeyEvent] 创建新帖子") + await self.client.create_note(content) await super().send(message) diff --git a/astrbot/core/platform/sources/misskey/misskey_utils.py b/astrbot/core/platform/sources/misskey/misskey_utils.py index 9a96b453..ebc95d8d 100644 --- a/astrbot/core/platform/sources/misskey/misskey_utils.py +++ b/astrbot/core/platform/sources/misskey/misskey_utils.py @@ -5,6 +5,68 @@ import astrbot.api.message_components as Comp from astrbot.api.platform import AstrBotMessage, MessageMember, MessageType +class FileIDExtractor: + """从 API 响应中提取文件 ID 的帮助类（无状态）。""" + + @staticmethod + def extract_file_id(result: Any) -> Optional[str]: + if not isinstance(result, dict): + return None + + id_paths = [ + lambda r: r.get("createdFile", {}).get("id"), + lambda r: r.get("file", {}).get("id"), + lambda r: r.get("id"), + ] + + for p in id_paths: + try: + if fid := p(result): + return fid + except Exception: + continue + + return None + + +class MessagePayloadBuilder: + """构建不同类型消息负载的帮助类（无状态）。""" + + @staticmethod + def build_chat_payload( + user_id: str, text: Optional[str], file_id: Optional[str] = None + ) -> Dict[str, Any]: + payload = {"toUserId": user_id} + if text: + payload["text"] = text + if file_id: + payload["fileId"] = file_id + return payload + + @staticmethod + def build_room_payload( + room_id: str, text: Optional[str], file_id: Optional[str] = None + ) -> Dict[str, Any]: + payload = {"toRoomId": room_id} + if text: + payload["text"] = text + if file_id: + payload["fileId"] = file_id + return payload + + @staticmethod + def build_note_payload( + text: Optional[str], file_ids: Optional[List[str]] = None, **kwargs + ) -> Dict[str, Any]: + payload: Dict[str, Any] = {} + if text: + payload["text"] = text + if file_ids: + payload["fileIds"] = file_ids + payload |= kwargs + return payload + + def serialize_message_chain(chain: List[Any]) -> Tuple[str, bool]: """将消息链序列化为文本字符串""" text_parts = [] @@ -15,11 +77,19 @@ def serialize_message_chain(chain: List[Any]) -> Tuple[str, bool]: if isinstance(component, Comp.Plain): return component.text elif isinstance(component, Comp.File): - file_name = getattr(component, "name", "文件") - return f"[文件: {file_name}]" + # 为文件组件返回占位符，但适配器仍会处理原组件 + return "[文件]" + elif isinstance(component, Comp.Image): + # 为图片组件返回占位符，但适配器仍会处理原组件 + return "[图片]" elif isinstance(component, Comp.At): has_at = True - return f"@{component.qq}" + # 优先使用name字段（用户名），如果没有则使用qq字段 + # 这样可以避免在Misskey中生成 @ 这样的无效提及 + if hasattr(component, "name") and component.name: + return f"@{component.name}" + else: + return f"@{component.qq}" elif hasattr(component, "text"): text = getattr(component, "text", "") if "@" in text: @@ -43,15 +113,22 @@ def serialize_message_chain(chain: List[Any]) -> Tuple[str, bool]: def resolve_message_visibility( - user_id: Optional[str], - user_cache: Dict[str, Any], - self_id: Optional[str], + user_id: Optional[str] = None, + user_cache: Optional[Dict[str, Any]] = None, + self_id: Optional[str] = None, + raw_message: Optional[Dict[str, Any]] = None, default_visibility: str = "public", ) -> Tuple[str, Optional[List[str]]]: - """解析 Misskey 消息的可见性设置""" + """解析 Misskey 消息的可见性设置 + + 可以从 user_cache 或 raw_message 中解析，支持两种调用方式： + 1. 基于 user_cache: resolve_message_visibility(user_id, user_cache, self_id) + 2. 基于 raw_message: resolve_message_visibility(raw_message=raw_message, self_id=self_id) + """ visibility = default_visibility visible_user_ids = None + # 优先从 user_cache 解析 if user_id and user_cache: user_info = user_cache.get(user_id) if user_info: @@ -66,38 +143,36 @@ def resolve_message_visibility( visible_user_ids = [uid for uid in visible_user_ids if uid] else: visibility = original_visibility + return visibility, visible_user_ids + + # 回退到从 raw_message 解析 + if raw_message: + original_visibility = raw_message.get("visibility", default_visibility) + if original_visibility == "specified": + visibility = "specified" + original_visible_users = raw_message.get("visibleUserIds", []) + sender_id = raw_message.get("userId", "") + + users_to_include = [] + if sender_id: + users_to_include.append(sender_id) + if self_id: + users_to_include.append(self_id) + + visible_user_ids = list(set(original_visible_users + users_to_include)) + visible_user_ids = [uid for uid in visible_user_ids if uid] + else: + visibility = original_visibility return visibility, visible_user_ids +# 保留旧函数名作为向后兼容的别名 def resolve_visibility_from_raw_message( raw_message: Dict[str, Any], self_id: Optional[str] = None ) -> Tuple[str, Optional[List[str]]]: - """从原始消息数据中解析可见性设置""" - visibility = "public" - visible_user_ids = None - - if not raw_message: - return visibility, visible_user_ids - - original_visibility = raw_message.get("visibility", "public") - if original_visibility == "specified": - visibility = "specified" - original_visible_users = raw_message.get("visibleUserIds", []) - sender_id = raw_message.get("userId", "") - - users_to_include = [] - if sender_id: - users_to_include.append(sender_id) - if self_id: - users_to_include.append(self_id) - - visible_user_ids = list(set(original_visible_users + users_to_include)) - visible_user_ids = [uid for uid in visible_user_ids if uid] - else: - visibility = original_visibility - - return visibility, visible_user_ids + """从原始消息数据中解析可见性设置（已弃用，使用 resolve_message_visibility 替代）""" + return resolve_message_visibility(raw_message=raw_message, self_id=self_id) def is_valid_user_session_id(session_id: Union[str, Any]) -> bool: @@ -128,6 +203,20 @@ def is_valid_room_session_id(session_id: Union[str, Any]) -> bool: ) +def is_valid_chat_session_id(session_id: Union[str, Any]) -> bool: + """检查 session_id 是否是有效的聊天 session_id (仅限chat%前缀)""" + if not isinstance(session_id, str) or "%" not in session_id: + return False + + parts = session_id.split("%") + return ( + len(parts) == 2 + and parts[0] == "chat" + and bool(parts[1]) + and parts[1] != "unknown" + ) + + def extract_user_id_from_session_id(session_id: str) -> str: """从 session_id 中提取用户 ID""" if "%" in session_id: @@ -149,21 +238,22 @@ def extract_room_id_from_session_id(session_id: str) -> str: def add_at_mention_if_needed( text: str, user_info: Optional[Dict[str, Any]], has_at: bool = False ) -> str: - """如果需要且没有@用户，则添加@用户""" + """如果需要且没有@用户，则添加@用户 + + 注意：仅在有有效的username时才添加@提及，避免使用用户ID + """ if has_at or not user_info: return text username = user_info.get("username") - nickname = user_info.get("nickname") + # 如果没有username，则不添加@提及，返回原文本 + # 这样可以避免生成 @ 这样的无效提及 + if not username: + return text - if username: - mention = f"@{username}" - if not text.startswith(mention): - text = f"{mention}\n{text}".strip() - elif nickname: - mention = f"@{nickname}" - if not text.startswith(mention): - text = f"{mention}\n{text}".strip() + mention = f"@{username}" + if not text.startswith(mention): + text = f"{mention}\n{text}".strip() return text @@ -197,6 +287,22 @@ def process_files( return file_parts +def format_poll(poll: Dict[str, Any]) -> str: + """将 Misskey 的 poll 对象格式化为可读字符串。""" + if not poll or not isinstance(poll, dict): + return "" + multiple = poll.get("multiple", False) + choices = poll.get("choices", []) + text_choices = [ + f"({idx}) {c.get('text', '')} [{c.get('votes', 0)}票]" + for idx, c in enumerate(choices, start=1) + ] + parts = ["[投票]", ("允许多选" if multiple else "单选")] + ( + ["选项: " + ", ".join(text_choices)] if text_choices else [] + ) + return " ".join(parts) + + def extract_sender_info( raw_data: Dict[str, Any], is_chat: bool = False ) -> Dict[str, Any]: @@ -248,7 +354,7 @@ def create_base_message( else: session_prefix = "note" session_id = f"{session_prefix}%{sender_info['sender_id']}" - message.type = MessageType.FRIEND_MESSAGE + message.type = MessageType.OTHER_MESSAGE message.session_id = ( session_id if sender_info["sender_id"] else f"{session_prefix}%unknown" @@ -303,6 +409,8 @@ def cache_user_info( "nickname": sender_info["nickname"], "visibility": raw_data.get("visibility", "public"), "visible_user_ids": raw_data.get("visibleUserIds", []), + # 保存原消息ID，用于回复时作为reply_id + "reply_to_note_id": raw_data.get("id"), } user_cache[sender_info["sender_id"]] = user_cache_data @@ -325,3 +433,106 @@ def cache_room_info( "visibility": "specified", "visible_user_ids": [client_self_id], } + + +async def resolve_component_url_or_path( + comp: Any, +) -> Tuple[Optional[str], Optional[str]]: + """尝试从组件解析可上传的远程 URL 或本地路径。 + + 返回 (url_candidate, local_path)。两者可能都为 None。 + 这个函数尽量不抛异常，调用方可按需处理 None。 + """ + url_candidate = None + local_path = None + + async def _get_str_value(coro_or_val): + """辅助函数：统一处理协程或普通值""" + try: + if hasattr(coro_or_val, "__await__"): + result = await coro_or_val + else: + result = coro_or_val + return result if isinstance(result, str) else None + except Exception: + return None + + try: + # 1. 尝试异步方法 + for method in ["convert_to_file_path", "get_file", "register_to_file_service"]: + if not hasattr(comp, method): + continue + try: + value = await _get_str_value(getattr(comp, method)()) + if value: + if value.startswith("http"): + url_candidate = value + break + else: + local_path = value + except Exception: + continue + + # 2. 尝试 get_file(True) 获取可直接访问的 URL + if not url_candidate and hasattr(comp, "get_file"): + try: + value = await _get_str_value(comp.get_file(True)) + if value and value.startswith("http"): + url_candidate = value + except Exception: + pass + + # 3. 回退到同步属性 + if not url_candidate and not local_path: + for attr in ("file", "url", "path", "src", "source"): + try: + value = getattr(comp, attr, None) + if value and isinstance(value, str): + if value.startswith("http"): + url_candidate = value + break + else: + local_path = value + break + except Exception: + continue + + except Exception: + pass + + return url_candidate, local_path + + +def summarize_component_for_log(comp: Any) -> Dict[str, Any]: + """生成适合日志的组件属性字典（尽量不抛异常）。""" + attrs = {} + for a in ("file", "url", "path", "src", "source", "name"): + try: + v = getattr(comp, a, None) + if v is not None: + attrs[a] = v + except Exception: + continue + return attrs + + +async def upload_local_with_retries( + api: Any, + local_path: str, + preferred_name: Optional[str], + folder_id: Optional[str], +) -> Optional[str]: + """尝试本地上传，返回 file id 或 None。如果文件类型不允许则直接失败。""" + try: + res = await api.upload_file(local_path, preferred_name, folder_id) + if isinstance(res, dict): + fid = res.get("id") or (res.get("raw") or {}).get("createdFile", {}).get( + "id" + ) + if fid: + return str(fid) + except Exception: + # 上传失败，直接返回 None，让上层处理错误 + return None + + return None diff --git a/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_server.py b/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_server.py index bc1e4b17..4a2eae74 100644 --- a/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_server.py +++ b/astrbot/core/platform/sources/qqofficial_webhook/qo_webhook_server.py @@ -15,12 +15,13 @@ class QQOfficialWebhook: self.appid = config["appid"] self.secret = config["secret"] self.port = config.get("port", 6196) + self.is_sandbox = config.get("is_sandbox", False) self.callback_server_host = config.get("callback_server_host", "0.0.0.0") if isinstance(self.port, str): self.port = int(self.port) - self.http: BotHttp = BotHttp(timeout=300) + self.http: BotHttp = BotHttp(timeout=300, is_sandbox=self.is_sandbox) self.api: BotAPI = BotAPI(http=self.http) self.token = Token(self.appid, self.secret) diff --git a/astrbot/core/platform/sources/satori/satori_adapter.py b/astrbot/core/platform/sources/satori/satori_adapter.py index efb75668..a3f4f53e 100644 --- a/astrbot/core/platform/sources/satori/satori_adapter.py +++ b/astrbot/core/platform/sources/satori/satori_adapter.py @@ -499,10 +499,36 @@ class SatoriPlatformAdapter(Platform): } return None + except ET.ParseError as e: + logger.warning(f"XML解析失败，使用正则提取: {e}") + return await self._extract_quote_with_regex(content) except Exception as e: logger.error(f"提取标签时发生错误: {e}") return None + async def _extract_quote_with_regex(self, content: str) -> Optional[dict]: + """使用正则表达式提取quote标签信息""" + import re + + quote_pattern = r"]*)>(.*?)" + match = re.search(quote_pattern, content, re.DOTALL) + + if not match: + return None + + attrs_str = match.group(1) + inner_content = match.group(2) + + id_match = re.search(r'id\s*=\s*["\']([^"\']*)["\']', attrs_str) + quote_id = id_match.group(1) if id_match else "" + content_without_quote = content.replace(match.group(0), "") + content_without_quote = content_without_quote.strip() + + return { + "quote": {"id": quote_id, "content": inner_content}, + "content_without_quote": content_without_quote, + } + async def _convert_quote_message(self, quote: dict) -> Optional[AstrBotMessage]: """转换引用消息""" try: @@ -574,7 +600,7 @@ class SatoriPlatformAdapter(Platform): root = ET.fromstring(processed_content) await self._parse_xml_node(root, elements) except ET.ParseError as e: - logger.error(f"解析 Satori 元素时发生解析错误: {e}, 错误内容: {content}") + logger.warning(f"解析 Satori 元素时发生解析错误: {e}, 错误内容: {content}") # 如果解析失败，将整个内容当作纯文本 if content.strip(): elements.append(Plain(text=content)) diff --git a/astrbot/core/platform/sources/satori/satori_event.py b/astrbot/core/platform/sources/satori/satori_event.py index f1a66da6..78325c9a 100644 --- a/astrbot/core/platform/sources/satori/satori_event.py +++ b/astrbot/core/platform/sources/satori/satori_event.py @@ -2,7 +2,18 @@ from typing import TYPE_CHECKING from astrbot.api import logger from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.platform import AstrBotMessage, PlatformMetadata -from astrbot.api.message_components import Plain, Image, At, File, Record, Video, Reply +from astrbot.api.message_components import ( + Plain, + Image, + At, + File, + Record, + Video, + Reply, + Forward, + Node, + Nodes, +) if TYPE_CHECKING: from .satori_adapter import SatoriPlatformAdapter @@ -48,55 +59,24 @@ class SatoriPlatformEvent(AstrMessageEvent): content_parts = [] for component in message.chain: - if isinstance(component, Plain): - text = ( - component.text.replace("&", "&") - .replace("<", "<") - .replace(">", ">") - ) - content_parts.append(text) + component_content = await cls._convert_component_to_satori_static( + component + ) + if component_content: + content_parts.append(component_content) - elif isinstance(component, At): - if component.qq: - content_parts.append(f'') - elif component.name: - content_parts.append(f'') + # 特殊处理 Node 和 Nodes 组件 + if isinstance(component, Node): + # 单个转发节点 + node_content = await cls._convert_node_to_satori_static(component) + if node_content: + content_parts.append(node_content) - elif isinstance(component, Image): - try: - image_base64 = await component.convert_to_base64() - if image_base64: - content_parts.append( - f'' - ) - except Exception as e: - logger.error(f"图片转换为base64失败: {e}") - - elif isinstance(component, File): - content_parts.append( - f'' - ) - - elif isinstance(component, Record): - try: - record_base64 = await component.convert_to_base64() - if record_base64: - content_parts.append( - f'