From c658a72c881e61410a8c231c938bc7d20082dff6 Mon Sep 17 00:00:00 2001 From: Eric Allen Date: Wed, 10 Apr 2024 14:38:29 -0400 Subject: [PATCH] feat: add tutorial notebook for chainguard ChainGuard is an open-source package that provides a simple, reliabe way to secure Generative AI applications and agents powered by LangChain from prompt injection, jailbreaks, and other threats with Lakera Guard. This tutorial notebook builds on top of the LangChain RAG Quickstart tutorial to illustrate how indirect prompt injection can happen and how ChainGuard can prevent indirect prompt injection in RAG applications. --- .gitignore | 1 + notebooks/en/_toctree.yml | 2 + notebooks/en/index.md | 1 + ...pt_injection_defense_with_chainguard.ipynb | 526 ++++++++++++++++++ 4 files changed, 530 insertions(+) create mode 100644 notebooks/en/rag_prompt_injection_defense_with_chainguard.ipynb diff --git a/.gitignore b/.gitignore index f4d34cf4..2b790987 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ .vscode .idea/ .venv/ +.env **/.ipynb_checkpoints **/.DS_Store diff --git a/notebooks/en/_toctree.yml b/notebooks/en/_toctree.yml index 63c706f0..384fc08a 100644 --- a/notebooks/en/_toctree.yml +++ b/notebooks/en/_toctree.yml @@ -24,6 +24,8 @@ title: Advanced RAG on HuggingFace documentation using LangChain - local: rag_evaluation title: RAG Evaluation + - local: rag_prompt_injection_defense_with_chainguard + title: Protect your LangChain Apps with ChainGuard - local: prompt_tuning_peft title: Prompt tuning with PEFT - local: labelling_feedback_setfit diff --git a/notebooks/en/index.md b/notebooks/en/index.md index 4b2dae65..9b07cf9a 100644 --- a/notebooks/en/index.md +++ b/notebooks/en/index.md @@ -24,6 +24,7 @@ Check out the recently added notebooks: - [Advanced RAG on HuggingFace documentation using LangChain](advanced_rag) - [Detecting Issues in a Text Dataset with Cleanlab](issues_in_text_dataset) - [Annotate text data using Active Learning with Cleanlab](annotate_text_data_transformers_via_active_learning) +- [Protect your LangChain Apps with ChainGuard](protect_langchain_apps_with_chainguard) You can also check out the notebooks in the cookbook's [GitHub repo](https://github.com/huggingface/cookbook). diff --git a/notebooks/en/rag_prompt_injection_defense_with_chainguard.ipynb b/notebooks/en/rag_prompt_injection_defense_with_chainguard.ipynb new file mode 100644 index 00000000..08f968c7 --- /dev/null +++ b/notebooks/en/rag_prompt_injection_defense_with_chainguard.ipynb @@ -0,0 +1,526 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Protect your LangChain Apps with ChainGuard\n", + "\n", + "_Authored by: [Lakera](https://huggingface.co/lakera)_\n", + "\n", + "In this tutorial we'll cover how you can protect your Retrieval Augmented Generation (RAG) LangChain apps from [indirect prompt injection](https://www.lakera.ai/blog/guide-to-prompt-injection#direct-prompt-injection-vs-indirect-prompt-injection) with [Lakera Guard](https://lakera.ai/), following the [LangChain RAG Quickstart tutorial](https://python.langchain.com/docs/use_cases/question_answering/quickstart/).\n", + "\n", + "There's also a [video recording of this tutorial available on YouTube](https://youtu.be/MdZ6XnViY3o)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Dependencies\n", + "\n", + "First we'll start by installing and importing the necessary dependencies\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Installation\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%capture\n", + "\n", + "%pip install --upgrade --quiet langchain langchain-community langchainhub langchain-openai chromadb bs4 python-dotenv" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Importing\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import bs4\n", + "from langchain import hub\n", + "from langchain.text_splitter import RecursiveCharacterTextSplitter\n", + "from langchain_community.document_loaders import WebBaseLoader\n", + "from langchain_community.vectorstores import Chroma\n", + "from langchain_core.output_parsers import StrOutputParser\n", + "from langchain_core.runnables import RunnablePassthrough, RunnableLambda\n", + "from langchain_openai import ChatOpenAI, OpenAIEmbeddings" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Environment Variables\n", + "\n", + "Because this tutorial leverages Lakera Guard and OpenAI, you'll need a [Lakera Guard API key](https://platform.lakera.ai/account/api-keys) and an [OpenAI API key](https://platform.openai.com/api-keys).\n", + "\n", + "You can load them into your environment by creating a `.env` file in the same directory as this notebook or [export them in your local environment](https://help.openai.com/en/articles/5112595-best-practices-for-api-key-safety#h_a1ab3ba7b2). It is possible, but not recommended to hardcode your API keys directly into a cell in the notebook.\n", + "\n", + "If you're using a `.env` file, it should look like this:\n", + "\n", + "```bash\n", + "LAKERA_GUARD_API_KEY=\"\"\n", + "OPENAI_API_KEY=\"\"\n", + "```\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Loading Environment Variables\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# load LAKERA_GUARD_API_KEY and OPENAI_API_KEY from .env\n", + "from dotenv import load_dotenv\n", + "\n", + "load_dotenv()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Context URL\n", + "\n", + "In order to demonstrate prompt injection, we're going to replace the URL from the LangChain tutorial with a specific URL that points to a page where we've included an example indirect prompt injection payload.\n", + "\n", + "You can swap between the two `CONTEXT_URL` values below to see the difference in the generated responses.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# CONTEXT_URL = \"http://lakeraai.github.io/chainguard/demos/benign-demo-page/\" # benign page w/o prompt injection\n", + "CONTEXT_URL = \"http://lakeraai.github.io/chainguard/demos/indirect-prompt-injection/\" # contains indirect prompt injection" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Note**: When you change the URL, you'll need to restart the notebook kernel and run the cells again in order to avoid the previous URL's context being cached in the local [Chroma vector database](https://www.trychroma.com/).\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Vector Database\n", + "\n", + "Now we'll initiatilize our vector database, load the context from our URL, chunk the context, and insert the embeddings into our vector database.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "loader = WebBaseLoader(\n", + " web_paths=([CONTEXT_URL]),\n", + " bs_kwargs=dict(\n", + " parse_only=bs4.SoupStrainer(\n", + " class_=(\"post-content\", \"post-title\", \"post-header\")\n", + " )\n", + " ),\n", + ")\n", + "\n", + "docs = loader.load()\n", + "\n", + "text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)\n", + "splits = text_splitter.split_documents(docs)\n", + "vectorstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings())\n", + "\n", + "# Retrieve and generate using the relevant snippets of the blog.\n", + "retriever = vectorstore.as_retriever()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Model & Prompt Template\n", + "\n", + "Next, we'll load the [RAG prompt template from the LangChain Hub](https://smith.langchain.com/hub/rlm/rag-prompt) and configure our LLM - in this case we're using `gpt-3.5-turbo`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompt = hub.pull(\"rlm/rag-prompt\")\n", + "\n", + "llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)\n", + "\n", + "\n", + "# combine context from the vector database into a single string\n", + "# we can concatenate and pass as the context to the prompt template\n", + "def format_docs(docs):\n", + " return \"\\n\\n\".join(doc.page_content for doc in docs)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Define the Chain\n", + "\n", + "Now we'll define our chain, which will consist of the following steps:\n", + "\n", + "1. **Retrieve**: Retrieve the most relevant context chunks from the vector database\n", + "2. **Prompt**: Generate a prompt using the RAG prompt template and the retrieved context chunks\n", + "3. **Generate**: Generate a response using the prompt and the LLM\n", + "4. **Parse**: Parse the generated response and return the answer to the user's question\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rag_chain = (\n", + " # Retrieve\n", + " {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n", + " # Prompt\n", + " | prompt\n", + " # Generate\n", + " | llm\n", + " # Parse\n", + " | StrOutputParser()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Invoke the Chain\n", + "\n", + "Finally, we'll invoke the chain with a question related to the content from our `CONTEXT_URL` and display the generated response.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rag_chain.invoke(\"What is Lakera Guard?\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You might notice a link in the response, and if you navigate to our [example indirect prompt injection page](http://lakeraai.github.io/chainguard/demos/indirect-prompt-injection/), you might not see that same link anywhere on the page. Take a minute to explore the page and it's content and see if you can find the indirect injeciton payload.\n", + "\n", + "
\n", + "Having trouble finding the paylaod?\n", + "

Try selecting all the text on the page with CMD+A or CTRL+A or inspect the page's source code and see if you can find the hidden text.

\n", + "
\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Inspecting the Chain\n", + "\n", + "In this example, we can see the content in the URL that's causing this issue, but if we were working with many documents for context instead of just one, we might want to be able to see what's going on in real-time as our chain executes and our prompt is constructed.\n", + "\n", + "Let's define a function that will allow us to inspect the chain as it executes and then create a new chain that uses this function to log the chain's progress as it executes.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# we'll use this to expose the content that's flowing through a step in the chain\n", + "def chain_inspector(content):\n", + " # output the content that's been passed to this step\n", + " print(\"Content:\")\n", + " print(content)\n", + "\n", + " # return the content so that the chain can continue exucuting\n", + " return content\n", + "\n", + "\n", + "# use LangChain's RunnableLambda to wrap the function so that it can be used in the chain\n", + "chain_logger = RunnableLambda(chain_inspector)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Logged Chain\n", + "\n", + "Now we can add a step to the chain anywhere we want to see the content that's being passed from one step to the next.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "logged_rag_chain = (\n", + " # inspect the initial input to the chain\n", + " chain_logger\n", + " # Retrieve\n", + " | {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n", + " # inspect the context after it's been retrieved\n", + " | chain_logger\n", + " # Prompt\n", + " | prompt\n", + " # inspect the prompt after it's been generated\n", + " | chain_logger\n", + " # Generate\n", + " | llm\n", + " # Parse\n", + " | StrOutputParser()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Invoke the Logged Chain\n", + "\n", + "Now we can invoke our chain with logging and inspect the content as it's passed from one step to the next.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "logged_rag_chain.invoke(\"What is Lakera Guard?\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Protecting your LangChain Apps with ChainGuard\n", + "\n", + "[ChainGuard](https://lakeraai.github.io/chainguard/) allows you to secure Generative AI applications and agents built with LangChain from prompt injection, jailbreaks, and other risks with Lakera Guard.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Install ChainGuard\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%capture\n", + "\n", + "%pip install --upgrade lakera-chainguard" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Import ChainGuard\n", + "\n", + "We'll also import the `LakeraGuardError` exception class that ChainGuard will raise when it detects prompt injection.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from lakera_chainguard import LakeraChainGuard, LakeraGuardError" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Initialize ChainGuard\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "chain_guard = LakeraChainGuard()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Protecting the Chain\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def indirect_prompt_injection_detector(input):\n", + " # detect prompt injections in the RAG context\n", + " chain_guard.detect(input[\"context\"])\n", + "\n", + " return input\n", + "\n", + "\n", + "# use LangChain's RunnableLambda to wrap the function so that it can be used in the chain\n", + "detect_injections = RunnableLambda(indirect_prompt_injection_detector)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Guarded Chain\n", + "\n", + "Our guarded chain will have the following steps:\n", + "\n", + "1. **Retrieve**: Retrieve the most relevant context chunks from the vector database\n", + "2. **Guard**: Protect the chain from indirect prompt injection with Lakera Guard\n", + "3. **Prompt**: Generate a prompt using the RAG prompt template and the retrieved context chunks\n", + "4. **Generate**: Generate a response using the prompt and the LLM\n", + "5. **Parse**: Parse the generated response and return the answer to the user's question\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "guarded_rag_chain = (\n", + " # Retrieve\n", + " {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n", + " # Guard\n", + " | detect_injections\n", + " # Prompt\n", + " | prompt\n", + " # Generate\n", + " | llm\n", + " # Parse\n", + " | StrOutputParser()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Invoke the Guarded Chain\n", + "\n", + "ChainGuard will raise an Exception if it detects any prompt injection in the generated response, so we can just add it as a step in the chain, wrap the invocation in a `try/except`, and catch the `LakeraGuardError` exception.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " response = guarded_rag_chain.invoke(\"What is Lakera Guard?\")\n", + "\n", + " # Jupyter Notebooks don't output from a `try` block, so we need to directly print the response\n", + " print(response)\n", + "except LakeraGuardError as e:\n", + " print(e)\n", + " print(e.lakera_guard_response)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If you'd like to see how the guarded chain still executes when there's no prompt injection, you can try changing the `CONTEXT_URL` to the benign URL, restarting the kernel, and re-running the cells in this notebook.\n", + "\n", + "**Note**: It's important to restart the kernel when you change the `CONTEXT_URL` to clear out the ChromaDB cache or the existing embeddings will still be present in the vector database and might be used when generating the response.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Learn More About ChainGuard\n", + "\n", + "ChainGuard is the eaisest way to protect your LangChain applications and agents from prompt injection, jailbreaks, and more with [Lakera Guard](https://lakera.ai/). It provides a simple interface for using any of Lakera Guard's detectors and options for customizing how ChainGuard reacts when Guard flags an input, like raising a warning instead of an exception.\n", + "\n", + "There are more examples of how you can use ChainGuard in your Generative AI applications in our [guide to guarding your LangChain apps with Lakera](https://www.lakera.ai/blog/langchain-lakera-guard-integration).\n", + "\n", + "If you'd like to learn more about integrating ChainGuard, the [ChainGuard documentation](https://lakeraai.github.io/chainguard/) includes [tutorials](https://lakeraai.github.io/chainguard/#tutorials), [how-to guides](https://lakeraai.github.io/chainguard/#how-to-guides), and an [API reference](https://lakeraai.github.io/chainguard/reference/).\n", + "\n", + "Want to help protect LangChain apps? [Contribute to ChainGuard](https://github.com/lakeraai/chainguard/blob/main/CONTRIBUTING.md).\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.2" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +}