This commit is contained in:
tt-P607
2025-08-21 00:05:39 +08:00
9 changed files with 1990 additions and 172 deletions

View File

@@ -1,119 +0,0 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
安全地向 llm_usage 表添加列 model_assign_name 及其索引
支持 MySQL / MariaDBPyMySQL 和 SQLitesqlite3
"""
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)

View File

@@ -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",

View File

@@ -63,4 +63,5 @@ Pillow
chromadb
asyncio
tavily-python
google-generativeai == 0.8.5
google-generativeai
cognee

View File

@@ -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. 检查用户是否被封禁

View File

@@ -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,

View File

@@ -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)

View File

@@ -1,3 +1,5 @@
# 有一个人想混点提交()
# 什么?混提交不带我一个喵~
import asyncio
import time
import signal

View File

@@ -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"""🛡️ 反注入系统状态报告

1991
uv.lock generated

File diff suppressed because it is too large Load Diff