Sejal-Siddapuram/PromptGuard

# PromptGuard 一个围绕 LLM 集成 pipeline 构建的多层 prompt injection 检测系统。它部署在 LLM (Ollama) 的前端，在对抗性输入到达模型之前对其进行拦截。 该项目作为 PESU Innovation Lab 研究实习的一部分，与 Akamai Technologies 合作开发。 ## 工作原理 每个传入的 prompt 都会经过三个层级的检测： **Tier 1 — Regex** 跨多个类别的模式匹配，包括指令覆盖、系统 prompt 提取尝试、越狱、编码技巧和 agent 操纵。速度快，零延迟，无需模型。 **Tier 2 — 语义搜索** 基于填充了已知攻击和良性 prompt 数据集的 ChromaDB 向量存储。在攻击向量和良性向量之间使用加权相似度投票，以计算 injection 概率得分。 **Tier 3 — LoRA 分类器** 经过 fine-tuning 的 Qwen3-0.6B 模型（`abedegno/prompt-injection-classifier-qwen3-0p6b`），用于对每个 prompt 进行不安全/安全分类。此项为可选 —— 如果被禁用或不可用，将平滑回退到双层 pipeline。 这三个得分将被合并为最终的风险得分。如果某个 prompt 超过了阈值，它会在 Ollama 收到之前被直接拦截。 ## 技术栈 - **FastAPI** — 检测 API 和中间件防护屏障 - **Ollama** (`qwen2.5:0.5b`) — 受保护的 LLM - **ChromaDB** — 用于语义搜索的向量存储 - **Streamlit** — 用于实时监控和分析的 SIEM 风格仪表板 - **sentence-transformers** (`all-MiniLM-L6-v2`) — 用于语义搜索的 embedding - **peft + transformers** — LoRA adapter 加载 一切均运行在 Docker 中。 ## 前置条件 - Docker + Docker Compose - 推荐 NVIDIA GPU（如果不可用，Ollama 将回退到 CPU） - 拥有读取 token 的 HuggingFace 账号 —— 这是下载 LoRA 模型和部分数据集所必需的 ## 设置 ### 1. 克隆仓库 ``` git clone https://github.com/Sejal-Siddapuram/PromptGuard.git cd PromptGuard ``` ### 2. HuggingFace token LoRA 分类器和部分数据集需要 HuggingFace 访问权限。请在 https://huggingface.co/settings/tokens 生成一个 token（只读权限即可）。 将其添加到你的 `docker-compose.yml` 中 `fastapi` 服务的 environment 下： ``` environment: - OLLAMA_URL=http://ollama:11434/api/generate - HF_TOKEN=your_token_here ``` 或者将其设置为机器上的环境变量并进行引用： ``` environment: - HF_TOKEN=${HF_TOKEN} ``` ### 3. 启动容器 ``` docker compose up --build -d ``` ### 4. 拉取 LLM ``` docker exec ollama ollama pull qwen2.5:0.5b ``` ### 5. 填充向量存储 此步骤将下载多个数据集（HackAPrompt、LLMail、Databricks Dolly、LMSYS Chat）并构建 ChromaDB 索引。首次运行需要几分钟。请在项目根目录下运行此命令 —— `chroma_db/` 文件夹已挂载到容器中，因此它会在重新构建后保持持久化。 ``` python populate_db.py ``` ## 用法 | 服务 | URL | | 仪表板 | http://localhost:8501 | | FastAPI | http://localhost:8000 | | Ollama | http://localhost:11434 | 发送聊天请求： ``` 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` 可禁用 Tier 3 并仅运行双层检测 | | `LORA_DEVICE` | `cpu` | `cpu` 或 `cuda` | | `LORA_THRESHOLD` | `0.10` | LoRA 不安全分类的概率截断值 | | `CHROMA_COLLECTION` | `prompts` | ChromaDB 集合名称 | | `HF_TOKEN` | — | 你的 HuggingFace 读取 token | ## 禁用 LoRA（低资源环境） 如果你在没有 GPU 的情况下运行，或者希望启动更快，请在你的 compose 文件中设置 `LORA_ENABLED=false`。系统将仅回退到 Tier 1 + Tier 2。检测仍然有效，只是没有了神经分类器层。 ## 代码修改后重新构建 ``` # 仅重新构建 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 ├── 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 ├── requirements.txt ├── requirements_dashboard.txt ├── Dockerfile └── docker-compose.yml ```

标签：AI安全网关, AI风险缓解, Kubernetes, LLM应用防火墙, LoRA微调, 向量数据库, 大语言模型安全, 提示注入防御, 机密管理, 源代码安全, 系统调用监控, 请求拦截, 逆向工具