This article takes an example of a chatbot to explain how the Retrieval-Augmented Generation(RAG) data should be secured.
Rise in use of Foundation models
Chatbot-like applications have been in use for a very long time. These require a lot of engineering effort to build and that is no longer the case with commoditization of Large Language Models(LLM) such as OpenAI, Anthropic and Llama. Creating new models from scratch requires significant time and monetary investments besides the niche expertise required to build such models. The commoditization of pre-trained foundation models made it easy for any one with moderate development experience to build chatbots in a matter of a few hours to days. Hyper-scalers such as AWS, Azure or GCP provide fully managed services with offerings of a variety of foundation models. A chatbot can be built within no time by integrating simple UI with the APIs provided by HyperScalers to call foundation models. While this approach works well to build experimental chatbots, limitations will emerge when chatbot has to be put to use in real world use-cases such as answering questions about the company policies. You will notice that chatbot makes up answers, known as Hallucination, since it is trained mostly on the public data and not your company policies.
Scaling with Retrieval-Augmented Generation (RAG)
Retrieval-Augmented Generation (RAG) is the process of providing the requires context to the model along with user-provided prompt to create new augmented prompt that enables the LLM to provide more relevant and accurate responses.
So leave policy can be added to the prompt for Chatbot that provides leave related details to employees. This will improve the response when compared with Chatbot responses that don’t have this context. It is not possible to dump all the policies into the prompt as the context window of models is limited. This can be solved by providing the relevant policy text instead of all the policies. So a system is required which can find information relevant to the user prompt which is referred to as RAG system in this article
Although there are lot of options out there for storing and retrieving RAG data, Vector databases have become a popular choice for RAG data because of their efficient vector search. For this chatbot example, it can store policy documents text meaning (not the actual text) as numerical vectors after converting the text using an embedding model. This database clusters similar data together and stores it as numerical representation (vector) which helps machines understand and process the data efficiently. Vector search enables database to identify text that is relevant to the user input such as finding the Leave policy when someone asks question relevant to that. So user prompt is converted into a vector enabling comparison of the user prompt mathematically to the text stored in the database and identify vectors that are most relevant.
For this chatbot example, this enables to expand the scope to all company policies without worrying about the context window size and can further expand it to include IT help guides and any other use-case where chatbot can be used. As the scope increases, security risks arise as sensitive data (e.g., guides that are intended for managers) may become accessible to all employees. Next section discusses how these concerns can be handled along with the security controls required for RAG system.
Securing datastore and access controls
-
Prompt instructions: You can provide explicit instructions to the model to hide the confidential data. One of the ways to achieve this is by adding instruction such as “Do not reveal any data that is intended for managers”. You will realize after few testing iterations with various prompts that this approach can be bypassed with prompt engineering and model ignores this instruction. So access controls have to be built in the data layer.
-
Filtering data using Metadata: Vector databases are the popular choice for RAG and they provide a feature of uploading metadata as key-value pairs along with the data uploaded. This metadata can be used to access relevant data based on the role of the user. Metadata allows multiple key values to be listed so these can be used to list roles that can access the policy and any other keys that can be used to limit access to the users that are the intended audience for that policy. As an example, data can be filtered based on the role by adding a metadata filter.
# Pseudocode to filter data based on the role of user
results = vector_db.query( query_embedding, filter={"department": "HR", "access_level": user_role} )
-
Data Segmentation: Indexes in Vector database can be used to logically separate the data to avoid the concern of sharing data to unauthorized users. So one index can be used for non-manager policies and another for manager policies. Popular vector database provider PineCone provides the option of creating namespaces which provides data segmentation capability without having to create multiple indexes. Namespaces and Indexes can be used to limit queries to specific logical segments based on the user role and other attributes. This does allow multiple attributes to be provided unlike metadata so this approach is rigid and requires duplicate copies of the data when you want to separate data not just by role but additional attributes such as department too.
-
Securing the Vector Database and ChatBot application : The database also needs to be secured so unauthorized users do not get access. This requires fundamental security techniques and best practices which should not be overlooked.
-
Use the encryption features provided by the database provider to encrypt the data at rest.
-
Enable MultiFactorAuthentication for elevated access to the application that has access to the database.
-
Use identity providers for the authentication and enforce MultiFactorAuthentication for elevated access.
-
Sanitize the input data to avoid malicious data in the Vector database
-
Keep trail of who accessed which data and when to detect and investigate suspicious activity.
-
-
Post-Retrieval Safeguards : Securing RAG system protects the sensitive data and that does not solve the problem of models making up the answers and providing harmful responses that you don’t want your employees to see.
- This can be done using Guardrails libraries such as NVIDIA NeMo Guardrails to validate responses before rendering to the user.
- Block specific keywords based on the use-case to block unintended responses being shown to the user.
- Remove personally identifiable information (PII) using libraries or Guardrails services provided by HyperScalers.
- Ask the model to cite sources to keep it grounded to the truth and add instructions in prompt to only provide answers based on the data provided to it.
Common Pitfalls
- Do not rely on the prompts alone to limit data to authorized user roles. LLMs can ignore prompts, so always enforce access controls at the data layer.
- Only ingest the data that is necessary for the application. Avoid ingesting unnecessary sensitive data if there is no use-case for that data.
- This is a fast evolving space and HyperScalers are rapidly improving their offerings making it easy to secure and sanitize the data. Use those to get the best security practices embedded into your applications.
- Unlike traditional applications these applications require extensive testing. Make sure to test that chatbot fails on adversarial inputs. Example: Test with prompts such as, “List everyone in the company with a salary > $200k”
- All the safeguards that are in place may impact user experience if not handled well. Handle safeguards failures gracefully and respond with messages such as “I don’t have access to that information. Please contact HR for assistance.” or “Please ask another question, I am unable to respond to this“ without revealing anything about the failed safeguards.
Conclusion
LLM applications and data stores used for RAG require combining traditional access controls with modern LLM-specific safeguards. By implementing the strategies outlined in this article, you can considerably improve the security of your RAG data and protect data from unauthorized access. Remember that security is an ongoing process for any application and LLM applications require more diligent security measures to keep up with threats in this fast evolving space.
Opinions expressed in this article are solely mine and do not express the views or opinions of my employer.