Refactor antiprompt flowchart and fix indentation
Updated the antipromptinjector flowchart for improved clarity and logic. Fixed indentation in expression_selector.py to correct code structure. Removed obsolete db_migration_plan.md documentation.
This commit is contained in:
雅诺狐
@@ -1,18 +1,18 @@
|
||||
```mermaid
|
||||
flowchart TD
|
||||
A([消息进入]) --> B{LLM反注入拦截
|
||||
是否启动？}
|
||||
B -- 否 --> Z([进入消息流])
|
||||
B -- 是 --> C{黑白名单判断}
|
||||
C -- 黑名单 --> D([丢弃消息])
|
||||
C -- 白名单 --> Z
|
||||
C -- 不在名单中 --> E{模式选择}
|
||||
E -- 仅正则表达模式 --> F[正则验证] --> G{验证通过？}
|
||||
G -- 是 --> Z
|
||||
G -- 否 --> D
|
||||
E -- 二次判定模式 --> H[正则验证] --> I{正则结果}
|
||||
I -- 正常 --> Z
|
||||
I -- 可疑 --> J[LLM二次判定] --> K{判定结果}
|
||||
K -- 安全 --> Z
|
||||
K -- 注入风险 --> D
|
||||
```
|
||||
A[消息进入系统] --> B{LLM反注入是否启动?}
|
||||
B -->|是| C{黑白名单检测}
|
||||
B -->|否| Y
|
||||
C -->|白名单| Y{继续进行消息处理}
|
||||
C -->|无记录| D{是否命中规则集}
|
||||
C -->|黑名单| X{丢弃消息}
|
||||
D -->|否| E{是否启动LLM二次分析}
|
||||
D -->|是| G{处理模式}
|
||||
E -->|是| F{提交LLM处理}
|
||||
E -->|否| Y
|
||||
F -->|LLM判定高危| G
|
||||
F -->|LLM判定无害| Y
|
||||
G -->|严格模式| X
|
||||
G -->|宽松模式| H{消息加盾}
|
||||
H --> Y
|
||||
```
|
||||
|
||||
@@ -140,47 +140,47 @@ class ExpressionSelector:
|
||||
(Expression.chat_id.in_(related_chat_ids)) & (Expression.type == "grammar")
|
||||
))
|
||||
|
||||
style_exprs = [
|
||||
{
|
||||
"situation": expr.situation,
|
||||
"style": expr.style,
|
||||
"count": expr.count,
|
||||
"last_active_time": expr.last_active_time,
|
||||
"source_id": expr.chat_id,
|
||||
"type": "style",
|
||||
"create_date": expr.create_date if expr.create_date is not None else expr.last_active_time,
|
||||
}
|
||||
for expr in style_query.scalars()
|
||||
]
|
||||
style_exprs = [
|
||||
{
|
||||
"situation": expr.situation,
|
||||
"style": expr.style,
|
||||
"count": expr.count,
|
||||
"last_active_time": expr.last_active_time,
|
||||
"source_id": expr.chat_id,
|
||||
"type": "style",
|
||||
"create_date": expr.create_date if expr.create_date is not None else expr.last_active_time,
|
||||
}
|
||||
for expr in style_query.scalars()
|
||||
]
|
||||
|
||||
grammar_exprs = [
|
||||
{
|
||||
"situation": expr.situation,
|
||||
"style": expr.style,
|
||||
"count": expr.count,
|
||||
"last_active_time": expr.last_active_time,
|
||||
"source_id": expr.chat_id,
|
||||
"type": "grammar",
|
||||
"create_date": expr.create_date if expr.create_date is not None else expr.last_active_time,
|
||||
}
|
||||
for expr in grammar_query.scalars()
|
||||
]
|
||||
grammar_exprs = [
|
||||
{
|
||||
"situation": expr.situation,
|
||||
"style": expr.style,
|
||||
"count": expr.count,
|
||||
"last_active_time": expr.last_active_time,
|
||||
"source_id": expr.chat_id,
|
||||
"type": "grammar",
|
||||
"create_date": expr.create_date if expr.create_date is not None else expr.last_active_time,
|
||||
}
|
||||
for expr in grammar_query.scalars()
|
||||
]
|
||||
|
||||
style_num = int(total_num * style_percentage)
|
||||
grammar_num = int(total_num * grammar_percentage)
|
||||
# 按权重抽样（使用count作为权重）
|
||||
if style_exprs:
|
||||
style_weights = [expr.get("count", 1) for expr in style_exprs]
|
||||
selected_style = weighted_sample(style_exprs, style_weights, style_num)
|
||||
else:
|
||||
selected_style = []
|
||||
if grammar_exprs:
|
||||
grammar_weights = [expr.get("count", 1) for expr in grammar_exprs]
|
||||
selected_grammar = weighted_sample(grammar_exprs, grammar_weights, grammar_num)
|
||||
else:
|
||||
selected_grammar = []
|
||||
|
||||
return selected_style, selected_grammar
|
||||
style_num = int(total_num * style_percentage)
|
||||
grammar_num = int(total_num * grammar_percentage)
|
||||
# 按权重抽样（使用count作为权重）
|
||||
if style_exprs:
|
||||
style_weights = [expr.get("count", 1) for expr in style_exprs]
|
||||
selected_style = weighted_sample(style_exprs, style_weights, style_num)
|
||||
else:
|
||||
selected_style = []
|
||||
if grammar_exprs:
|
||||
grammar_weights = [expr.get("count", 1) for expr in grammar_exprs]
|
||||
selected_grammar = weighted_sample(grammar_exprs, grammar_weights, grammar_num)
|
||||
else:
|
||||
selected_grammar = []
|
||||
|
||||
return selected_style, selected_grammar
|
||||
|
||||
def update_expressions_count_batch(self, expressions_to_update: List[Dict[str, Any]], increment: float = 0.1):
|
||||
"""对一批表达方式更新count值，按chat_id+type分组后一次性写入数据库"""
|
||||
|
||||
@@ -1,138 +0,0 @@
|
||||
# 数据库自动迁移模块 (`db_migration.py`) 设计文档
|
||||
|
||||
## 1. 目标
|
||||
|
||||
创建一个自动化的数据库迁移模块，用于在应用启动时检查数据库结构，并自动进行以下修复：
|
||||
|
||||
1. **创建缺失的表**：如果代码模型中定义的表在数据库中不存在，则自动创建。
|
||||
2. **添加缺失的列**：如果数据库中的某个表现有的列比代码模型中定义的少，则自动添加缺失的列。
|
||||
|
||||
## 2. 实现思路
|
||||
|
||||
我们将使用 SQLAlchemy 的 `Inspector` 来获取数据库的元数据（即实际的表和列信息），并将其与 `SQLAlchemy` 模型（`Base.metadata`）中定义的结构进行比较。
|
||||
|
||||
核心逻辑分为以下几个步骤：
|
||||
|
||||
1. **获取数据库引擎**：从现有代码中获取已初始化的 SQLAlchemy 引擎实例。
|
||||
2. **获取 Inspector**：通过引擎创建一个 `Inspector` 对象。
|
||||
3. **获取所有模型**：遍历 `Base.metadata.tables`，获取所有在代码中定义的表模型。
|
||||
4. **获取数据库中所有表名**：使用 `inspector.get_table_names()` 获取数据库中实际存在的所有表名。
|
||||
5. **创建缺失的表**：通过比较模型表名和数据库表名，找出所有缺失的表，并使用 `table.create(engine)` 来创建它们。
|
||||
6. **检查并添加缺失的列**：
|
||||
* 遍历每一个代码中定义的表模型。
|
||||
* 使用 `inspector.get_columns(table_name)` 获取数据库中该表的实际列。
|
||||
* 比较模型列和实际列，找出所有缺失的列。
|
||||
* 对于每一个缺失的列，生成一个 `ALTER TABLE ... ADD COLUMN ...` 的 SQL 语句，并执行它。
|
||||
|
||||
## 3. 伪代码实现
|
||||
|
||||
```python
|
||||
# mmc/src/common/database/db_migration.py
|
||||
|
||||
from sqlalchemy import inspect, text
|
||||
from sqlalchemy.engine import Engine
|
||||
from src.common.database.sqlalchemy_models import Base, get_engine
|
||||
from src.common.logger import get_logger
|
||||
|
||||
logger = get_logger("db_migration")
|
||||
|
||||
def check_and_migrate_database():
|
||||
"""
|
||||
检查数据库结构并自动迁移（添加缺失的表和列）。
|
||||
"""
|
||||
logger.info("正在检查数据库结构并执行自动迁移...")
|
||||
engine = get_engine()
|
||||
inspector = inspect(engine)
|
||||
|
||||
# 1. 获取数据库中所有已存在的表名
|
||||
db_table_names = set(inspector.get_table_names())
|
||||
|
||||
# 2. 遍历所有在代码中定义的模型
|
||||
for table_name, table in Base.metadata.tables.items():
|
||||
logger.debug(f"正在检查表: {table_name}")
|
||||
|
||||
# 3. 如果表不存在，则创建它
|
||||
if table_name not in db_table_names:
|
||||
logger.info(f"表 '{table_name}' 不存在，正在创建...")
|
||||
try:
|
||||
table.create(engine)
|
||||
logger.info(f"表 '{table_name}' 创建成功。")
|
||||
except Exception as e:
|
||||
logger.error(f"创建表 '{table_name}' 失败: {e}")
|
||||
continue
|
||||
|
||||
# 4. 如果表已存在，则检查并添加缺失的列
|
||||
db_columns = {col['name'] for col in inspector.get_columns(table_name)}
|
||||
model_columns = {col.name for col in table.c}
|
||||
|
||||
missing_columns = model_columns - db_columns
|
||||
if not missing_columns:
|
||||
logger.debug(f"表 '{table_name}' 结构一致，无需修改。")
|
||||
continue
|
||||
|
||||
logger.info(f"在表 '{table_name}' 中发现缺失的列: {', '.join(missing_columns)}")
|
||||
with engine.connect() as connection:
|
||||
for column_name in missing_columns:
|
||||
column = table.c[column_name]
|
||||
|
||||
# 构造并执行 ALTER TABLE 语句
|
||||
# 注意：这里的实现需要考虑不同数据库（SQLite, MySQL）的语法差异
|
||||
# 为了简化，我们先使用一个通用的格式，后续可以根据需要进行扩展
|
||||
try:
|
||||
column_type = column.type.compile(engine.dialect)
|
||||
sql = f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}"
|
||||
|
||||
# 可以在这里添加对默认值、非空约束等的处理
|
||||
|
||||
connection.execute(text(sql))
|
||||
logger.info(f"成功向表 '{table_name}' 添加列 '{column_name}'。")
|
||||
except Exception as e:
|
||||
logger.error(f"向表 '{table_name}' 添加列 '{column_name}' 失败: {e}")
|
||||
|
||||
# 提交事务
|
||||
if connection.in_transaction():
|
||||
connection.commit()
|
||||
|
||||
logger.info("数据库结构检查与自动迁移完成。")
|
||||
|
||||
```
|
||||
|
||||
## 4. 集成到启动流程
|
||||
|
||||
为了让这个迁移模块在应用启动时自动运行，我们需要在 `mmc/src/common/database/sqlalchemy_models.py` 的 `initialize_database` 函数中调用它。
|
||||
|
||||
修改后的 `initialize_database` 函数将如下所示：
|
||||
|
||||
```python
|
||||
# mmc/src/common/database/sqlalchemy_models.py
|
||||
|
||||
# ... (其他 import)
|
||||
from src.common.database.db_migration import check_and_migrate_database # 导入新函数
|
||||
|
||||
# ... (代码)
|
||||
|
||||
def initialize_database():
|
||||
"""初始化数据库引擎和会话"""
|
||||
global _engine, _SessionLocal
|
||||
|
||||
if _engine is not None:
|
||||
return _engine, _SessionLocal
|
||||
|
||||
# ... (数据库连接和引擎创建逻辑)
|
||||
|
||||
_engine = create_engine(database_url, **engine_kwargs)
|
||||
_SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=_engine)
|
||||
|
||||
# 在这里，我们不再直接调用 create_all
|
||||
# Base.metadata.create_all(bind=_engine)
|
||||
|
||||
# 而是调用新的迁移函数，它会处理表的创建和列的添加
|
||||
check_and_migrate_database()
|
||||
|
||||
logger.info(f"SQLAlchemy数据库初始化成功: {config.database_type}")
|
||||
return _engine, _SessionLocal
|
||||
|
||||
# ... (其他代码)
|
||||
```
|
||||
|
||||
通过这样的修改，我们就可以在不改变现有初始化流程入口的情况下，无缝地集成自动化的数据库结构检查和修复功能。
|
||||
Reference in New Issue
Block a user