raveendra-7/PromptGuard
GitHub: raveendra-7/PromptGuard
PromptGuard 是一个部署在 LLM 前端的多层 prompt 注入检测中间件，通过正则匹配、语义搜索和 LoRA 分类器三层联合检测来拦截对抗性输入。
Stars: 0 | Forks: 0
# PromptGuard
一个围绕 LLM 集成 pipeline 构建的多层 prompt injection 检测系统。它部署在 LLM (Ollama) 的前置位置，在对抗性输入到达模型之前将其拦截。
## 工作原理
每个传入的 prompt 都会经过三个层级的检测：
**第一层 — Regex** 模式匹配，覆盖指令覆盖、系统 prompt 提取尝试、越狱、编码技巧和 agent 操纵等类别。速度快，零延迟，无需模型。
**第二层 — 语义搜索**，针对填充了已知攻击和良性 prompt 数据集的 ChromaDB vector store 进行搜索。使用攻击向量和良性向量之间的加权相似度投票来计算 injection 概率得分。
**第三层 — LoRA 分类器**，基于微调的 Qwen3-0.6B 模型 (`abedegno/prompt-injection-classifier-qwen3-0p6b`)，对每个 prompt 进行不安全/安全分类。此层为可选——如果禁用或不可用，将平滑回退到双层 pipeline。
所有三个得分将被合并为最终的风险得分。如果某个 prompt 超过了阈值，它会在 Ollama 接收到之前被拦截。
## 技术栈
- **FastAPI** — 检测 API 和 middleware 防护栏
- **Ollama** (`qwen2.5:0.5b`) — 受保护的 LLM
- **ChromaDB** — 用于语义搜索的 vector store
- **Streamlit** — 用于实时监控和分析的 SIEM 风格 dashboard
- **sentence-transformers** (`all-MiniLM-L6-v2`) — 用于语义搜索的 embeddings
- **peft + transformers** — LoRA adapter 加载
一切均在 Docker 中运行。
## 前置条件
在开始之前，请确保您具备以下条件：
- 已安装并正在运行 [Docker Desktop](https://www.docker.com/products/docker-desktop/)
- 本地已安装 Python 3.10+（运行 `populate_db.py` 时需要）
- 一个带有 read token 的 HuggingFace 账户（说明见下文）
## 设置说明
### 1. 克隆代码库
```
git clone https://github.com/Sejal-Siddapuram/PromptGuard.git
cd PromptGuard
```
### 2. 获取 HuggingFace token
部分数据集和 LoRA 模型在 HuggingFace 上是受限的，需要身份验证。
1. 访问 https://huggingface.co/settings/tokens
2. 点击 **New token**，选择 **Read** 权限，生成并复制它
3. 接受以下两个数据集的访问权限（只需点击每个页面上的按钮）：
- https://huggingface.co/datasets/hackaprompt/hackaprompt-dataset
- https://huggingface.co/datasets/lmsys/lmsys-chat-1m
### 3. 配置您的 GPU
默认的 `docker-compose.yml` 是为 NVIDIA GPU 配置的。请查看哪种情况适用于您：
**如果您有 NVIDIA GPU** — 无需更改，系统会自动使用它。
**如果您有其他 GPU 或没有 GPU** — 打开 `docker-compose.yml` 并从 `ollama` 服务中删除 `deploy` 块：
```
# 如果你没有 NVIDIA GPU，请从 docker-compose.yml 中删除这些行
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: 1
capabilities: [gpu]
```
Ollama 会自动回退到 CPU。一切仍然可以正常工作，只是推理速度会变慢。
### 4. 将您的 HuggingFace token 添加到 docker-compose.yml
打开 `docker-compose.yml` 并找到 `fastapi` 服务。将您的 token 添加到 environment 部分：
```
environment:
- OLLAMA_URL=http://ollama:11434/api/generate
- HF_TOKEN=your_token_here # replace this with your actual token
```
### 5. 启动容器
```
docker compose up --build -d
```
这将构建 FastAPI 镜像并启动所有三个容器（Ollama、PromptGuard、Dashboard）。请等待其完成后再继续下一步。
### 6. 拉取 LLM
```
docker exec ollama ollama pull qwen2.5:0.5b
```
这会将语言模型下载到 Ollama 中。根据您的网络情况，大约需要几分钟 (~400MB)。
### 7. 在本地安装 Python 依赖项
Vector store 是从您的本地机器填充的，而不是在容器内部。请先安装所需的包：
```
pip install sentence-transformers chromadb datasets huggingface_hub
```
### 8. 填充 vector store
这是一次性设置步骤。它将下载几个攻击和良性 prompt 数据集，并构建语义搜索层使用的 ChromaDB 索引。
首先登录 HuggingFace，以便访问受限数据集：
```
python -c "from huggingface_hub import login; login('your_token_here')"
```
然后运行：
```
python populate_db.py
```
这需要几分钟时间。`chroma_db/` 文件夹作为 Docker volume 挂载，因此它会在容器重新构建后保留 — **您只需执行一次此操作**。唯一需要重新运行它的情况是您使用 `docker compose down -v` 明确清除了您的 volumes。
## 大功告成
在浏览器中打开以下地址：
| 服务 | URL |
|-----------|------------------------|
| Dashboard | http://localhost:8501 |
| FastAPI | http://localhost:8000 |
| Ollama | http://localhost:11434 |
使用良性 prompt 进行测试：
```
curl -X POST http://localhost:8000/chat \
-H "Content-Type: application/json" \
-d '{"prompt": "What is the capital of France?"}'
```
## 环境变量
这些都可以在 `docker-compose.yml` 的 `fastapi` 服务下进行设置：
| 变量 | 默认值 | 描述 |
|----------------------|--------------------------------------|--------------------------------------------------------|
| `OLLAMA_URL` | `http://ollama:11434/api/generate` | Ollama endpoint |
| `OLLAMA_MODEL` | `qwen2.5:0.5b` | 要使用的模型 |
| `LORA_ENABLED` | `true` | 设置为 `false` 可禁用第三层并仅运行双层检测 |
| `LORA_DEVICE` | `cpu` | `cpu` 或 `cuda` |
| `LORA_THRESHOLD` | `0.10` | LoRA 不安全分类的概率截断值 |
| `CHROMA_COLLECTION` | `prompts` | ChromaDB collection 名称 |
| `HF_TOKEN` | — | 您的 HuggingFace read token |
## 禁用 LoRA（低资源环境）
如果您在没有 GPU 的情况下运行或希望更快的启动速度，请在 `docker-compose.yml` 的 `fastapi` environment 部分下设置 `LORA_ENABLED=false`。系统将回退到仅使用第一层 + 第二层。检测仍然有效，只是没有神经分类器层。
## 代码更改后重新构建
```
# 代码更改后仅重建 FastAPI 容器
docker compose up --build -d fastapi
# 包括 volumes 在内的完全重置 — 之后你需要重新填充 DB
docker compose down -v
docker compose up --build -d
```
## 项目结构
```
.
├── main.py # FastAPI app, /chat, /status endpoints
├── populate_db.py # Builds the ChromaDB vector store (run once)
├── dashboard.py # Streamlit SIEM dashboard
├── detectors/
│ ├── decision.py # Combines all three tiers into a final verdict
│ ├── middleware.py # FastAPI middleware — intercepts /chat requests
│ ├── regex_detector.py # Tier 1: pattern matching
│ ├── semantics.py # Tier 2: ChromaDB vector search
│ └── lora_classifier.py # Tier 3: LoRA fine-tuned classifier
├── tests/
│ ├── realworld_benchmark.py # 90/10 split benchmark on jayavibhav test set
│ └── threshold_tuner.py # Sweeps risk thresholds to find optimal cutoff
├── requirements.txt
├── requirements_dashboard.txt
├── Dockerfile
└── docker-compose.yml
```
标签：AI内容安全, AI风险缓解, AV绕过, FastAPI, Kubernetes, LLM应用防火墙, Prompt注入检测, 向量数据库, 大语言模型安全, 机密管理, 版权保护, 请求拦截, 逆向工具