Joyaljose0/LLM-prompt-injection-detector
GitHub: Joyaljose0/LLM-prompt-injection-detector
基于微调 BERT 模型的 LLM 提示词注入检测系统，在恶意指令到达 AI Agent 之前进行拦截和评分。
Stars: 0 | Forks: 0
# LLM 提示词注入检测器
本项目构建了一个自主分类器，旨在检测并拦截 AI agent pipeline 中的恶意提示词注入。随着企业部署处理外部不可信数据（如网页、文档和工具输出）的 LLM agent，它们极易受到隐藏在这些内容中的对抗性指令的攻击。该系统会拦截输入，对其恶意意图进行评分，并在其到达执行 agent 之前将其拦截。
## 项目架构
该架构由三个主要阶段组成：
1. **外部内容获取**：模拟来自网页、电子邮件或用户提示词的输入。
2. **注入检测器**：利用序列分类模型（例如，经过微调的 BERT 模型）来分配注入概率分数的系统。
3. **决策与路由**：
- 如果分数超过阈值（`> 0.5`），则该提示词将被**拦截并记录**，以供分析师审查。
- 如果安全，则将该提示词传递给 **LLM Agent** 进行正常执行。
## 功能
- **FastAPI 后端**：用于处理和评分文本的高性能 REST API。
- **检测器模型包装器**：包含一个备用的启发式检测器，以及一个 HuggingFace BERT 微调脚本（`train.py`），用于在合成注入对上训练模型。
- **LangChain 模拟 Agent**：一个演示安全执行路径的模拟 agent。
- **高级前端演示**：采用 Vanilla JS/HTML/CSS 构建、具有动态 UI 动画并映射出系统架构的视觉效果极佳的暗色模式仪表板。
## 快速开始
### 前置条件
- Python 3.8+
- Node.js（可选，仅当您希望通过开发服务器提供静态前端服务时才需要，否则只需在浏览器中打开 `index.html` 即可）
### 安装说明
1. 克隆仓库：
git clone https://github.com/Joyaljose0/LLM-prompt-injection-detector.git
cd LLM-prompt-injection-detector
2. 安装后端依赖：
pip install -r backend/requirements.txt
### 运行应用程序
1. **启动 FastAPI 后端**：
uvicorn backend.api:app --reload
API 将在 `http://127.0.0.1:8000` 提供。
2. **打开前端演示**：
只需在您的 Web 浏览器中打开 `frontend/index.html` 即可。
### 训练模型（可选）
要在提供的合成数据集上训练 `bert-tiny` 模型：
```
python backend/train.py
```
这将训练模型并将其保存到 `backend/injection_model/`。检测器将在后端重启时自动加载它。
## 技术栈
- **AI/ML**：PyTorch、HuggingFace Transformers、scikit-learn
- **后端**：FastAPI、Uvicorn、Pydantic、LangChain
- **前端**：HTML5、Vanilla CSS（Glassmorphism 设计）、Vanilla JavaScript
## 许可证
MIT
标签：凭据扫描, 数据可视化, 逆向工具