Merge branch 'master' of https://github.com/MaiBot-Plus/MaiMbot-Pro-Max
This commit is contained in:
@@ -1,119 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
安全地向 llm_usage 表添加列 model_assign_name 及其索引
|
||||
支持 MySQL / MariaDB(PyMySQL) 和 SQLite(sqlite3)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
import sqlite3
|
||||
import pymysql
|
||||
from contextlib import closing
|
||||
|
||||
# ===== 按需修改 =====
|
||||
# 1) MySQL 示例
|
||||
# CONNECTION = dict(
|
||||
# host="127.0.0.1",
|
||||
# port=3306,
|
||||
# user="root",
|
||||
# password="123456",
|
||||
# database="test_db",
|
||||
# charset="utf8mb4",
|
||||
# driver="mysql" # 关键标识
|
||||
# )
|
||||
|
||||
# 2) SQLite 示例
|
||||
CONNECTION = dict(
|
||||
database="test.db", # 数据库文件路径
|
||||
driver="sqlite"
|
||||
)
|
||||
# =====================
|
||||
|
||||
COLUMN_NAME = "model_assign_name"
|
||||
COLUMN_TYPE = "VARCHAR(100)"
|
||||
INDEX_NAME = f"idx_llmusage_{COLUMN_NAME}"
|
||||
TABLE_NAME = "llm_usage"
|
||||
|
||||
|
||||
# --------------------------------------------------
|
||||
# MySQL 实现
|
||||
# --------------------------------------------------
|
||||
def _mysql_ensure(cfg: dict) -> None:
|
||||
with closing(pymysql.connect(**{k: v for k, v in cfg.items() if k != "driver"})) as conn, closing(conn.cursor()) as cur:
|
||||
# 1. 列是否存在
|
||||
cur.execute("""
|
||||
SELECT COUNT(*) FROM information_schema.COLUMNS
|
||||
WHERE table_schema = %s
|
||||
AND table_name = %s
|
||||
AND column_name = %s
|
||||
""", (cfg["database"], TABLE_NAME, COLUMN_NAME))
|
||||
col_exists = cur.fetchone()[0] > 0
|
||||
|
||||
# 2. 索引是否存在
|
||||
cur.execute("""
|
||||
SELECT COUNT(*) FROM information_schema.STATISTICS
|
||||
WHERE table_schema = %s
|
||||
AND table_name = %s
|
||||
AND index_name = %s
|
||||
""", (cfg["database"], TABLE_NAME, INDEX_NAME))
|
||||
idx_exists = cur.fetchone()[0] > 0
|
||||
|
||||
if not col_exists:
|
||||
sql = f"ALTER TABLE {TABLE_NAME} ADD COLUMN {COLUMN_NAME} {COLUMN_TYPE}"
|
||||
print(f"[MySQL DDL] {sql}")
|
||||
cur.execute(sql)
|
||||
|
||||
if not idx_exists:
|
||||
sql = f"CREATE INDEX {INDEX_NAME} ON {TABLE_NAME} ({COLUMN_NAME})"
|
||||
print(f"[MySQL DDL] {sql}")
|
||||
cur.execute(sql)
|
||||
|
||||
conn.commit()
|
||||
|
||||
|
||||
# --------------------------------------------------
|
||||
# SQLite 实现
|
||||
# --------------------------------------------------
|
||||
def _sqlite_ensure(cfg: dict) -> None:
|
||||
db = cfg["database"]
|
||||
with closing(sqlite3.connect(db)) as conn, closing(conn.cursor()) as cur:
|
||||
# 1. 列是否存在
|
||||
cur.execute("PRAGMA table_info({})".format(TABLE_NAME))
|
||||
cols = {row[1] for row in cur.fetchall()}
|
||||
col_exists = COLUMN_NAME in cols
|
||||
|
||||
# 2. 索引是否存在
|
||||
cur.execute("PRAGMA index_list({})".format(TABLE_NAME))
|
||||
idxs = {row[1] for row in cur.fetchall()}
|
||||
idx_exists = INDEX_NAME in idxs
|
||||
|
||||
# SQLite ≥3.35 支持 ALTER TABLE ADD COLUMN IF NOT EXISTS
|
||||
if not col_exists:
|
||||
sql = f"ALTER TABLE {TABLE_NAME} ADD COLUMN {COLUMN_NAME} {COLUMN_TYPE}"
|
||||
print(f"[SQLite DDL] {sql}")
|
||||
cur.execute(sql)
|
||||
|
||||
if not idx_exists:
|
||||
sql = f"CREATE INDEX {INDEX_NAME} ON {TABLE_NAME} ({COLUMN_NAME})"
|
||||
print(f"[SQLite DDL] {sql}")
|
||||
cur.execute(sql)
|
||||
|
||||
conn.commit()
|
||||
|
||||
|
||||
# --------------------------------------------------
|
||||
# 调度器
|
||||
# --------------------------------------------------
|
||||
def ensure_column_and_index(cfg: dict) -> None:
|
||||
driver = cfg.get("driver", "").lower()
|
||||
if driver == "mysql":
|
||||
_mysql_ensure(cfg)
|
||||
elif driver == "sqlite":
|
||||
_sqlite_ensure(cfg)
|
||||
else:
|
||||
raise ValueError("connection.driver 必须是 'mysql' 或 'sqlite'")
|
||||
print("✅ 完成。")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ensure_column_and_index(CONNECTION)
|
||||
@@ -10,6 +10,7 @@ dependencies = [
|
||||
"asyncddgs>=0.1.0a1",
|
||||
"asyncio>=4.0.0",
|
||||
"beautifulsoup4>=4.13.4",
|
||||
"cognee>=0.2.3",
|
||||
"colorama>=0.4.6",
|
||||
"cryptography>=45.0.5",
|
||||
"customtkinter>=5.2.2",
|
||||
@@ -24,6 +25,7 @@ dependencies = [
|
||||
"json-repair>=0.47.6",
|
||||
"json5>=0.12.1",
|
||||
"jsonlines>=4.0.0",
|
||||
"langfuse==2.46.2",
|
||||
"lxml>=6.0.0",
|
||||
"maim-message>=0.3.8",
|
||||
"matplotlib>=3.10.3",
|
||||
@@ -31,10 +33,11 @@ dependencies = [
|
||||
"numpy>=2.2.6",
|
||||
"openai>=1.95.0",
|
||||
"opencv-python>=4.11.0.86",
|
||||
"packaging>=25.0",
|
||||
"packaging==24.1",
|
||||
"pandas>=2.3.1",
|
||||
"peewee>=3.18.2",
|
||||
"pillow>=11.3.0",
|
||||
"pip-check-reqs>=2.5.5",
|
||||
"psutil>=7.0.0",
|
||||
"pyarrow>=20.0.0",
|
||||
"pydantic>=2.11.7",
|
||||
|
||||
@@ -63,4 +63,5 @@ Pillow
|
||||
chromadb
|
||||
asyncio
|
||||
tavily-python
|
||||
google-generativeai == 0.8.5
|
||||
google-generativeai
|
||||
cognee
|
||||
@@ -60,11 +60,12 @@ class AntiPromptInjector:
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 统计更新
|
||||
await self.statistics.update_stats(total_messages=1)
|
||||
# 1. 检查系统是否启用
|
||||
if not self.config.enabled:
|
||||
return ProcessResult.ALLOWED, None, "反注入系统未启用"
|
||||
|
||||
# 统计更新 - 只有在系统启用时才进行统计
|
||||
await self.statistics.update_stats(total_messages=1)
|
||||
logger.debug(f"开始处理消息: {message.processed_plain_text}")
|
||||
|
||||
# 2. 检查用户是否被封禁
|
||||
|
||||
@@ -10,6 +10,7 @@ from typing import Dict, Any, Optional
|
||||
|
||||
from src.common.logger import get_logger
|
||||
from src.common.database.sqlalchemy_models import AntiInjectionStats, get_db_session
|
||||
from src.config.config import global_config
|
||||
|
||||
logger = get_logger("anti_injector.statistics")
|
||||
|
||||
@@ -19,7 +20,8 @@ class AntiInjectionStatistics:
|
||||
|
||||
def __init__(self):
|
||||
"""初始化统计管理器"""
|
||||
pass
|
||||
self.session_start_time = datetime.datetime.now()
|
||||
"""当前会话开始时间"""
|
||||
|
||||
async def get_or_create_stats(self):
|
||||
"""获取或创建统计记录"""
|
||||
@@ -78,6 +80,22 @@ class AntiInjectionStatistics:
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
try:
|
||||
# 检查反注入系统是否启用
|
||||
if not global_config.anti_prompt_injection.enabled:
|
||||
return {
|
||||
"status": "disabled",
|
||||
"message": "反注入系统未启用",
|
||||
"uptime": "N/A",
|
||||
"total_messages": 0,
|
||||
"detected_injections": 0,
|
||||
"blocked_messages": 0,
|
||||
"shielded_messages": 0,
|
||||
"detection_rate": "N/A",
|
||||
"average_processing_time": "N/A",
|
||||
"last_processing_time": "N/A",
|
||||
"error_count": 0
|
||||
}
|
||||
|
||||
stats = await self.get_or_create_stats()
|
||||
|
||||
# 计算派生统计信息 - 处理None值
|
||||
@@ -88,10 +106,13 @@ class AntiInjectionStatistics:
|
||||
detection_rate = (detected_injections / total_messages * 100) if total_messages > 0 else 0
|
||||
avg_processing_time = (processing_time_total / total_messages) if total_messages > 0 else 0
|
||||
|
||||
# 使用当前会话开始时间计算运行时间,而不是数据库中的start_time
|
||||
# 这样可以避免重启后显示错误的运行时间
|
||||
current_time = datetime.datetime.now()
|
||||
uptime = current_time - stats.start_time
|
||||
uptime = current_time - self.session_start_time
|
||||
|
||||
return {
|
||||
"status": "enabled",
|
||||
"uptime": str(uptime),
|
||||
"total_messages": total_messages,
|
||||
"detected_injections": detected_injections,
|
||||
|
||||
@@ -151,7 +151,7 @@ class Prompt(str):
|
||||
|
||||
@staticmethod
|
||||
def _process_escaped_braces(template) -> str:
|
||||
"""处理模板中的转义花括号,将 \{ 和 \} 替换为临时标记""" # type: ignore
|
||||
"""处理模板中的转义花括号,将 \\{ 和 \\} 替换为临时标记""" # type: ignore
|
||||
# 如果传入的是列表,将其转换为字符串
|
||||
if isinstance(template, list):
|
||||
template = "\n".join(str(item) for item in template)
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
# 有一个人想混点提交()
|
||||
# 什么?混提交不带我一个喵~
|
||||
import asyncio
|
||||
import time
|
||||
import signal
|
||||
|
||||
@@ -32,8 +32,14 @@ class AntiInjectorStatusCommand(BaseCommand):
|
||||
anti_injector = get_anti_injector()
|
||||
stats = await anti_injector.get_stats()
|
||||
|
||||
if stats.get("stats_disabled"):
|
||||
return True, "反注入系统统计功能已禁用", True
|
||||
# 检查反注入系统是否禁用
|
||||
if stats.get("status") == "disabled":
|
||||
await self.send_text("❌ 反注入系统未启用\n\n💡 请在配置文件中启用反注入功能后重试")
|
||||
return True, "反注入系统未启用", True
|
||||
|
||||
if stats.get("error"):
|
||||
await self.send_text(f"❌ 获取状态失败: {stats['error']}")
|
||||
return False, f"获取状态失败: {stats['error']}", True
|
||||
|
||||
status_text = f"""🛡️ 反注入系统状态报告
|
||||
|
||||
|
||||
Reference in New Issue
Block a user