diff --git a/src/chat/antipromptinjector/流程图.md b/src/chat/antipromptinjector/流程图.md index d1d0e0b0b..2cc9533a3 100644 --- a/src/chat/antipromptinjector/流程图.md +++ b/src/chat/antipromptinjector/流程图.md @@ -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 -``` \ No newline at end of file + 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 +``` diff --git a/src/chat/express/expression_selector.py b/src/chat/express/expression_selector.py index 61b5d48b9..69b81174d 100644 --- a/src/chat/express/expression_selector.py +++ b/src/chat/express/expression_selector.py @@ -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分组后一次性写入数据库""" diff --git a/src/common/database/db_migration_plan.md b/src/common/database/db_migration_plan.md deleted file mode 100644 index 591a94f88..000000000 --- a/src/common/database/db_migration_plan.md +++ /dev/null @@ -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 - -# ... (其他代码) -``` - -通过这样的修改，我们就可以在不改变现有初始化流程入口的情况下，无缝地集成自动化的数据库结构检查和修复功能。 \ No newline at end of file