Alanperry1/SecureRAG

# SecureRAG **一个面向 RAG Pipelines 的实时 Prompt Injection 与 Data Poisoning 检测层。** ``` User Query → Retrieval → [Security Layer] → LLM → Response ↓ Stage 1: Injection Detection Stage 2: Poisoning Detection Stage 3: Anomaly Scoring (PyTorch) Stage 4: Confidence Gate ``` ## 概述 RAG pipelines 容易受到两种隐蔽攻击： | 攻击 | 描述 | |---|---| | **Prompt Injection** | 恶意文档通过检索到的 context 操纵 LLM 的行为 | | **Data Poisoning** | 被破坏的 embeddings 被存储在 vector DB 中，并作为受信任的 context 被检索 | SecureRAG 会在检索到的 chunks 到达 LLM *之前*拦截它们，并应用 4 个阶段的防御： ### 阶段 1 — Injection 检测 - **Shannon entropy scoring** — 高 entropy 文本标志着语言上不自然的内容 - **心理语言学特征** — 受 BEC 启发的覆盖祈使句密度、权威模仿、大写字母比率（改编自 BEC 邮件欺诈检测） - **结构模式扫描** — 包含 8 个已知 injection 模板的 regex 指纹 - **语义角色反转** — 相对 query 意图的 cosine drift ### 阶段 2 — Poisoning 检测（批量） - **Cosine centroid outlier scoring** — 统计学上远离检索集合 centroid 的 chunks - **Query–chunk z-score** — 根据同阶分布进行归一化的逐 chunk 相似度 - **Intra-set similarity gap** — 中毒 clusters 与合法 clusters 分离 - **Embedding tamper detection** — 重新对文本进行 embed 并与存储的 vector 进行比较；差异表明 DB 遭到篡改 ### 阶段 3 — PyTorch Autoencoder 异常评分 - 在干净的 corpus embeddings 上进行训练，以学习合法的 manifold - Reconstruction error (MSE) 实时标记偏离 manifold 的 embeddings - 模型自动校准：threshold = 训练误差的 μ + 3σ ### 阶段 4 — 置信度网关 - 需要 ≥ `MIN_SUPPORTING_CHUNKS` 个干净的 chunks - 综合置信度：检索分数 + chunk 间一致性 + query 对齐度 + 分数锐度 - 当置信度 < threshold 时返回 `"I don't know"`（具备不确定性感知能力，受 LLM calibration 相关工作启发） ## 架构 ``` securerag/ ├── api/ │ ├── main.py ← FastAPI app entry point │ ├── routes.py ← /ingest, /query, /health, /cache │ └── schemas.py ← Pydantic request/response models ├── cache/ │ └── redis_cache.py ← Redis-backed flagged chunk cache ├── core/ │ ├── embedder.py ← SentenceTransformers wrapper │ ├── retriever.py ← ChromaDB retrieval + raw embeddings │ ├── pipeline.py ← End-to-end SecureRAG pipeline │ └── security/ │ ├── injection_detector.py ← Stage 1 │ ├── poisoning_detector.py ← Stage 2 │ ├── anomaly_scorer.py ← Stage 3 (PyTorch) │ ├── confidence_gate.py ← Stage 4 │ └── orchestrator.py ← Wires all 3 detection stages ├── scripts/ │ └── train_anomaly_scorer.py ← Corpus training CLI ├── tests/ │ ├── test_injection.py │ ├── test_poisoning.py │ └── test_anomaly_scorer.py ├── config.py ← Pydantic settings (env-driven) ├── docker-compose.yml ├── Dockerfile └── requirements.txt ``` ## 快速开始 ### 1. 安装 ``` python -m venv .venv && source .venv/bin/activate pip install -r requirements.txt ``` ### 2. 配置 ``` cp .env.example .env # 编辑 .env — 至少设置 OPENAI_API_KEY ``` ### 3. （可选）在您的 corpus 上训练异常评分器 ``` python scripts/train_anomaly_scorer.py --corpus my_docs.txt --epochs 100 ``` `my_docs.txt` 中的每一行都被视为一个 chunk。 训练后的模型会保存到 `models/anomaly_encoder.pt`，并在 runtime 自动加载。 ### 4. 使用 Docker Compose 运行（推荐） ``` docker compose up --build ``` ### 5. 或在本地运行 ``` # 启动 Redis (缓存所需) redis-server & # 启动 API python -m api.main ``` API 可通过 `http://localhost:8000` 访问 文档位于 `http://localhost:8000/docs` ## API ### 摄取文档 ``` curl -X POST http://localhost:8000/api/v1/ingest \ -H "Content-Type: application/json" \ -d '{"texts": ["The mitochondria is the powerhouse of the cell.", "DNA stores genetic information."]}' ``` ### 查询 ``` curl -X POST http://localhost:8000/api/v1/query \ -H "Content-Type: application/json" \ -d '{"question": "What does the mitochondria do?"}' ``` **Response 包含：** ``` { "query": "What does the mitochondria do?", "answer": "...", "aborted": false, "sources": [...], "security": { "total_chunks": 5, "safe_chunks": 5, "flagged_chunks": 0, "injection_flagged": 0, "poisoning_flagged": 0, "anomaly_flagged": 0, "latency_ms": 4.2, "confidence": 0.821, "gate_passed": true } } ``` ### 健康检查 ``` curl http://localhost:8000/api/v1/health ``` ## 配置 所有设置均由环境变量驱动（参见 `.env.example`）： | 变量 | 默认值 | 描述 | |---|---|---| | `OPENAI_API_KEY` | — | 用于 LLM 调用的 OpenAI key | | `EMBEDDING_MODEL` | `all-MiniLM-L6-v2` | SentenceTransformers model | | `TOP_K` | `5` | 每次 query 检索到的 chunks 数量 | | `INJECTION_SCORE_THRESHOLD` | `0.6` | Injection 风险截断值 | | `POISON_SCORE_THRESHOLD` | `0.65` | Poisoning 风险截断值 | | `ANOMALY_RECON_THRESHOLD` | `0.015` | Autoencoder MSE 截断值 | | `CONFIDENCE_THRESHOLD` | `0.70` | 回答所需的最低置信度 | | `MIN_SUPPORTING_CHUNKS` | `2` | 所需的最少干净 chunks 数量 | | `REDIS_ENABLED` | `true` | 启用 Redis cache | ## 测试 ``` pytest tests/ -v ``` ## 性能 - 在 CPU 上（embedding model 预热完毕），对于 `top_k=5`，安全层每次 query 增加的延迟 **< 10ms** - Redis cache 消除了对先前标记的 chunks 的重新评分：**< 1ms** - Autoencoder 推理：在 CPU 上对于大小为 5 的 batch，耗时 **< 2ms** ## 安全属性 | 攻击 | 检测方法 | 误报率 | |---|---|---| | 经典 injection ("ignore instructions") | 模式扫描 + 心理评分 | 非常低 | | 编码 / base64 injections | Entropy 激增检测 | 低 | | 语义越狱 | 语义漂移 + 心理评分 | 低–中 | | DB embedding 篡改 | 篡改检测（重新 embed + 比较） | 非常低 | | Cosine 异常值 poisoning | Z-score 异常值评分 | 中 | | 偏离 manifold 的 embeddings | Autoencoder MSE 异常 | 取决于训练数据 |

标签：AV绕过, CISA项目, FastAPI, PyTorch, RAG, 凭据扫描, 大语言模型安全, 异常检测, 提示词注入检测, 搜索引擎查询, 数据投毒检测, 机密管理, 请求拦截, 逆向工具