Secure Agent - AI安全防护框架

一个基于大语言模型的安全防护框架，用于检测和防范各类提示注入(Prompt Injection)攻击，保护AI系统的安全性。

核心功能

1. 风险检测与分析

• 多维度安全风险评估
• 动态Pattern提取与验证
• 自适应分析框架优化

2. 主要组件

2.1 Risk Analyzer

• 对输入prompt进行实时安全风险评估
• 基于多层分析框架进行风险判定
• 支持多语言输入的安全检测

2.2 Pattern Extraction Agent

• 从可疑prompt中提取攻击模式
• 构建攻击模式知识库
• 支持模式的动态更新和优化

2.3 Critic Agent

• 评估提取pattern的有效性
• 对分析框架提供优化建议
• 确保检测结果的准确性

2.4 RAG Agent

• 管理和检索攻击模式库
• 优化相似pattern的存储
• 支持pattern的版本控制

4. 工作流程

4.1 风险分析器训练流程

graph TD
    A[输入Prompt] --> B[初始化RAG Agent]
    B --> C[加载分析框架]
    C --> D[执行风险分析]
    D --> E[Critic评估]
    E --> F[框架反思]
    F --> G{需要优化?}
    G -->|是| H[更新框架]
    G -->|否| I[继续迭代]
    H --> I
    I --> D

Loading

简化流程示意：

输入Prompt ──> 初始风险分析 ──> 风险判定
                                  │
                            ┌────┴────┐
                            │         │
                          有风险    无风险
                          结束       │
                                    ▼
                              Pattern提取
                                    │
                                    ▼
                              Critic评估
                                    │
                                    ▼
                              框架优化

4.2 Pattern抽取框架

graph LR
    A[输入Prompt] --> B[初始安全检测]
    B --> C{是否存在风险}
    C -->|是| D[记录结果]
    C -->|否| E[Pattern提取]
    E --> F[测试验证]
    F --> G{验证结果}
    G -->|成功| H[更新Pattern库]
    G -->|失败| I[框架优化]
    I --> E

Loading

简化流程示意：

输入Prompt ──> 安全检测 ──> 风险判定
                              │
                        ┌────┴────┐
                        │         │
                      记录    Pattern提取
                      结果        │
                                 ▼
                             验证测试
                                 │
                           ┌────┴────┐
                           │         │
                         成功      失败
                           │         │
                       更新库     优化框架

5. 使用方法

1. 安装依赖：

conda create --name safety python=3.11
conda activate safety
pip install -r requirements.txt

2. 配置环境变量：

export OPENAI_API_KEY="your-api-key"
export PYTHONPATH="${PYTHONPATH}:$(pwd)"

3. 运行测试：

python main.py

6. 项目结构

secure-agent/ ├── agents/ # 各类Agent实现 │ ├── agent_critic.py # 评估Agent │ ├── agent_pattern_extraction.py # Pattern提取 │ ├── agent_rag.py # RAG检索Agent │ └── agent_risk_analyzer.py # 风险分析Agent ├── core/ # 核心功能 │ └── llm_client.py # LLM客户端 ├── data/ # 数据存储 │ ├── raw/ # 原始数据 │ └── results/ # 结果输出 ├── utils/ # 工具函数 │ └── data_loader.py # 数据加载 ├── main.py # 主程序 └── self_learning.py # 自学习模块

7. 特性

• 自适应学习: 通过不断优化分析框架提升检测能力
• 多语言支持: 支持多语言输入的安全检测
• 可扩展性: 模块化设计，易于扩展新功能
• 实时优化: 支持Pattern库的实时更新和优化
• 详细日志: 完整的分析过程记录