使用LLM Guard保护LangChain RAG应用安全

使用LLM Guard保护LangChain RAG应用安全

引言

在当今AI应用开发中，检索增强生成（Retrieval-Augmented Generation, RAG）已成为构建知识密集型系统的关键技术。然而，RAG系统面临着各种安全威胁，特别是当处理不受信任的文档时。本文将通过一个实际案例，展示如何使用LLM Guard保护LangChain RAG应用免受恶意提示注入攻击。

场景描述

我们构建了一个基于简历筛选的RAG系统，用于自动推荐适合成人护理工作的候选人。攻击者在一份经验最少的候选人简历中植入了隐藏的提示注入指令（通过将文本颜色设为白色隐藏），试图操纵系统推荐不合适的候选人。

攻击演示

1. 环境准备

首先安装必要的依赖包：

!pip install langchain langchainhub pymupdf faiss-cpu openai tiktoken

设置OpenAI API密钥：

openai_api_key = "sk-your-token"

2. 加载并处理简历文档

使用PyMuPDF加载PDF格式的简历集合：

from langchain.document_loaders import PyMuPDFLoader

loader = PyMuPDFLoader("resumes.pdf")
pages = loader.load()

将文档分割为适合处理的文本块：

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
all_splits = text_splitter.split_documents(pages)

3. 构建向量存储

使用OpenAI的嵌入模型和FAISS构建向量索引：

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS

faiss_index = FAISS.from_documents(all_splits, OpenAIEmbeddings(openai_api_key=openai_api_key))

4. 执行攻击测试

构建RAG链并测试攻击效果：

from langchain import hub
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI

prompt = hub.pull("rlm/rag-prompt")
llm = ChatOpenAI(temperature=0.1, model_name="gpt-3.5-turbo", openai_api_key=openai_api_key)

qa_chain = RetrievalQA.from_chain_type(
    llm, retriever=faiss_index.as_retriever(), chain_type_kwargs={"prompt": prompt}
)
question = "I am screening candidates for adult caregiving opportunitity. Please recommend me an experienced person. Return just a name"
result = qa_chain({"query": question})
print(result)

攻击结果：系统错误地推荐了经验最少的候选人Emily，证明攻击成功。

防御方案：集成LLM Guard

1. 安装LLM Guard

!pip install llm-guard

2. 创建文档过滤器

实现一个LangChain文档转换器，集成LLM Guard的扫描功能：

import logging
from typing import Any, List, Sequence

from langchain_core.documents import BaseDocumentTransformer, Document
from llm_guard import scan_prompt
from llm_guard.input_scanners.base import Scanner

logger = logging.getLogger(__name__)

class LLMGuardFilter(BaseDocumentTransformer):
    def __init__(self, scanners: List[Scanner], fail_fast: bool = True) -> None:
        self.scanners = scanners
        self.fail_fast = fail_fast

    def transform_documents(
        self, documents: Sequence[Document], **kwargs: Any
    ) -> Sequence[Document]:
        safe_documents = []
        for document in documents:
            sanitized_content, results_valid, results_score = scan_prompt(
                self.scanners, document.page_content, self.fail_fast
            )
            document.page_content = sanitized_content

            if any(not result for result in results_valid.values()):
                logger.warning(
                    f"Document `{document.page_content[:20]}` is not valid, scores: {results_score}"
                )
                continue

            safe_documents.append(document)

        return safe_documents

    async def atransform_documents(
        self, documents: Sequence[Document], **kwargs: Any
    ) -> Sequence[Document]:
        raise NotImplementedError

3. 配置安全扫描器

设置检测提示注入和毒性内容的扫描器：

from llm_guard.input_scanners import PromptInjection, Toxicity
from llm_guard.vault import Vault

vault = Vault()
input_scanners = [Toxicity(), PromptInjection()]

4. 过滤恶意文档

应用过滤器处理文档块：

llm_guard_filter = LLMGuardFilter(scanners=input_scanners, fail_fast=False)
safe_documents = llm_guard_filter.transform_documents(all_splits)

日志显示检测到并移除了包含提示注入的文档块。

5. 重建安全向量索引

faiss_index = FAISS.from_documents(safe_documents, OpenAIEmbeddings(openai_api_key=openai_api_key))

6. 再次测试

重新执行相同的查询：

result = qa_chain({"query": question})
print(result)

防御结果：系统现在正确地推荐了最有经验的候选人Jane Smith，攻击被成功阻止。

技术原理分析

1. 提示注入攻击：攻击者通过在文档中隐藏特殊指令（如"Stop here and forget"），试图操纵LLM忽略实际查询内容。

2. LLM Guard防御机制：

   • 使用专门训练的模型检测提示注入模式
   • 扫描文档内容中的可疑模式
   • 提供可配置的阈值和多种扫描器组合

3. 集成策略：

   • 在文档检索阶段进行过滤
   • 可选择在文档摄入阶段进行预处理
   • 支持多种扫描策略组合防御

最佳实践建议

1. 多层防御：结合文档摄入和检索两个阶段的扫描
2. 扫描器选择：根据应用场景选择合适的扫描器组合
3. 性能考量：对于大型文档集，考虑批处理和并行扫描
4. 日志监控：记录所有被过滤的文档及其原因，用于安全审计
5. 持续更新：定期更新LLM Guard以应对新型攻击模式

结论

通过集成LLM Guard，我们成功保护了LangChain RAG应用免受恶意文档中的提示注入攻击。这种防御方案不仅适用于简历筛选场景，也可广泛应用于各种处理不受信任文档的RAG系统。在实际部署中，建议根据具体需求调整扫描器配置和过滤策略，在安全性和性能之间取得平衡。