gardel/Mofox-Core
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修复handle_message_storage 中模式与结果的逻辑不一致

Browse Source
This commit is contained in:
LuiKlee
2025-10-08 17:57:53 +08:00
committed by Windpicker-owo
parent 69db860a1f
commit 3f75c06237
1 changed files with 30 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
src/chat/antipromptinjector/anti_injector.py
View File

@@ -221,33 +221,41 @@ class AntiPromptInjector:
self, result: ProcessResult, modified_content: str | None, reason: str, message_data: dict self, result: ProcessResult, modified_content: str | None, reason: str, message_data: dict
) -> None: ) -> None:
"""处理违禁消息的数据库存储，根据处理模式决定如何处理""" """处理违禁消息的数据库存储，根据处理模式决定如何处理"""
if result == ProcessResult.BLOCKED_INJECTION or result == ProcessResult.COUNTER_ATTACK: mode = self.config.process_mode
# 严格模式和反击模式：删除违禁消息记录 message_id = message_data.get("message_id")
if self.config.process_mode in ["strict", "counter_attack"]:
await self._delete_message_from_storage(message_data)
logger.info(f"[{self.config.process_mode}模式] 违禁消息已从数据库中删除: {reason}")
elif result == ProcessResult.SHIELDED: if not message_id:
# 宽松模式：替换消息内容为加盾版本 logger.warning("无法处理消息存储：缺少 message_id")
if modified_content and self.config.process_mode == "lenient": return
# 更新消息数据中的内容
message_data["processed_plain_text"] = modified_content
message_data["raw_message"] = modified_content
await self._update_message_in_storage(message_data, modified_content)
logger.info(f"[宽松模式] 违禁消息内容已替换为加盾版本: {reason}")
elif result in [ProcessResult.BLOCKED_INJECTION, ProcessResult.SHIELDED] and self.config.process_mode == "auto": if mode == "strict":
# 自动模式：根据威胁等级决定 if result == ProcessResult.BLOCKED_INJECTION:
await self._delete_message_from_storage(message_data)
logger.info(f"[严格模式] 违禁消息已从数据库中删除: {reason}")
elif result == ProcessResult.SHIELDED:
if modified_content:
await self._update_message_in_storage(message_data, modified_content)
logger.info(f"[严格模式] 违禁消息内容已替换为加盾版本: {reason}")
elif mode == "lenient":
if result == ProcessResult.SHIELDED:
if modified_content:
await self._update_message_in_storage(message_data, modified_content)
logger.info(f"[宽松模式] 违禁消息内容已替换为加盾版本: {reason}")
elif mode == "auto":
if result == ProcessResult.BLOCKED_INJECTION: if result == ProcessResult.BLOCKED_INJECTION:
# 高威胁：删除记录
await self._delete_message_from_storage(message_data) await self._delete_message_from_storage(message_data)
logger.info(f"[自动模式] 高威胁消息已删除: {reason}") logger.info(f"[自动模式] 高威胁消息已删除: {reason}")
elif result == ProcessResult.SHIELDED and modified_content: elif result == ProcessResult.SHIELDED:
# 中等威胁：替换内容 if modified_content:
message_data["processed_plain_text"] = modified_content await self._update_message_in_storage(message_data, modified_content)
message_data["raw_message"] = modified_content logger.info(f"[自动模式] 中等威胁消息已加盾: {reason}")
await self._update_message_in_storage(message_data, modified_content)
logger.info(f"[自动模式] 中等威胁消息已加盾: {reason}") elif mode == "counter_attack":
if result == ProcessResult.COUNTER_ATTACK:
await self._delete_message_from_storage(message_data)
logger.info(f"[反击模式] 违禁消息已从数据库中删除: {reason}")
@staticmethod @staticmethod
async def _delete_message_from_storage(message_data: dict) -> None: async def _delete_message_from_storage(message_data: dict) -> None:
@@ -265,7 +273,6 @@ class AntiPromptInjector:
async with get_db_session() as session: async with get_db_session() as session:
# 删除对应的消息记录 # 删除对应的消息记录
stmt = delete(Messages).where(Messages.message_id == message_id) stmt = delete(Messages).where(Messages.message_id == message_id)
# 注意: 异步会话需要 await 执行，否则 result 是 coroutine，无法获取 rowcount
result = await session.execute(stmt) result = await session.execute(stmt)
await session.commit() await session.commit()