Secure your RAG

With the increasing shift towards digitization, businesses and organizations constantly explore new ways to improve their operations and customer experience. A leading approach in this space involves the development of advanced customer support bots, internal knowledge graphs, and Q&A systems.

Retrieval Augmented Generation (RAG) applications are becoming more prevalent to make these more efficient. RAG’s ability to blend pre-trained models with proprietary data is a game-changer. However, ensuring that these applications are used safely without compromising data integrity is crucial. In this blog, we discuss the critical elements of RAG and introduce a way to manage RAG workflows with security and governance.

RAG: An Overview

Retrieval Augmented Generation, commonly known as RAG, is a cutting-edge architecture in artificial intelligence that blends the capabilities of large-scale neural language models with information retrieval systems. Instead of generating responses based solely on pre-trained knowledge, RAG retrieves relevant documents or data fragments from vast datasets. Then, it utilizes these fragments as context for the subsequent text generation. This allows RAG to tap into specific, up-to-date, or domain-specific knowledge, making it particularly effective for applications such as chatbots and Q&A systems that require real-time information access and customization based on external data.

Security and Controls

Javelin LLM Gateway is an Enterprise-grade LLM (Large Language Model) Gateway that enables enterprises to apply policy controls, adhere to governance measures, and enforce comprehensive security guardrails, including data leak prevention, to ensure safe and compliant model use. Javelin is tailored to empower Operational teams, enabling them to manage LLM-enabled applications and oversee their model access effectively. This blog shows how we can effectively harden RAG workflows with security, governance, and privacy controls.

Setting up a Secure RAG Workflow

Setting up a Retrieval Augmented Generation (RAG) system involves combining a retrieval mechanism with a generation model. Below is a high-level overview of the process:

1. Data Collection and Processing:

• Document Dataset: Gather a comprehensive dataset, typically consisting of documents, articles, or any information the model might need to refer to. Typically, this may be internal corporate documents, knowledge bases, policies, contract documents, etc.,

• Pre-process Data: Clean and format the data to ensure it’s consistent for indexing and retrieval. This may involve deciding what data you want to index and what data you want to leave out (e.g., sensitive or confidential data).

Security Tip: This is the phase where it is generally easy for sensitive or secure documents to make their way into your RAG Document Dataset inadvertently. Two key issues need to be managed:

• Sensitive Sources — Before selecting large drives, directories, or mounts, ensure you are not mixing sensitive data into RAG systems.

• Data Provenance — make sure you can track where data comes from. Uncertainty around sources often introduces vulnerabilities or inaccuracies.

2. Create an Embedding Space:

• Encode Documents: Use an encoder (often a transformer-based model) to convert each document in the dataset into dense vector representations. While setting up a simple workflow for a handful of documents is easy, the task becomes more complex as the data volume grows.

• There are a variety of embedding models, each tailored for specific data types. For instance, long-form documents with long sentences, like literature, require different embedding strategies from short sentences like chat, while software code requires a different consideration. Chunking the document data into smaller pieces also follows similar considerations.

Secure Your Model Access Using Javelin

Javelin is an LLM Gateway best suited to simply provide a single model access point from within your RAG workflows. Routes to various models may be securely provisioned on the gateway — this prevents your RAG application from having to understand, maintain, and keep track of models. This central model management also makes it easy to quickly select newer models as they become available or switch to different embedding models from different applications.

It’s easy to install the Javelin Python SDK:

‍

‍

Now, you can easily create one or more routes. Each route represents an endpoint that can be centrally managed with security and governance guardrails:

‍

‍

Once you have provisioned one or more routes, replace the model URLs in your apps and instead have them use Javelin.

Love Langchain? We love it, too! Here is an easy way (~1 line of code) to do this from your Langchain-enabled apps:

‍

‍

3. Store Embeddings:

Store these vector representations in a vector database, allowing for efficient nearest-neighbor lookups.

Security Tip: Three key security issues need to be carefully managed in this phase:

• Sensitive Information Leaks — Original data can sometimes be reverse-engineered or inferred from embeddings. So always start with non-sensitive data and ensure embeddings don’t leak.

• Access Control — ensure that only authorized personnel can access and generate embeddings, especially when dealing with sensitive or proprietary data.

• Membership Inference Attacks — given a specific embedding, attackers may be able to infer that a specific data point was part of the document set.

Secure Model Access With Javelin — with Javelin, your model access is always secure. Since all LLM access goes through Javelin, you have complete visibility of the applications using various models, how much they use them, and what data is being passed to models.

Budget Guardrails: You can even set up cost guardrails to control budgets for large document RAG jobs!

‍

‍

Fine-tune the Generator Model (Optional):

• While only sometimes necessary, if you have specific domain knowledge or need the model to have a particular style, you can fine-tune a pre-trained model on a custom dataset.

Security Tip: While security for fine-tuning is a complete article in itself, the key things to keep in mind at this point are:

• Model Confidentiality — if the base model is proprietary (e.g., Open AI Models), you need to ensure that the fine-tuned model doesn’t expose details about its training.

• Operational Security & Data Infrastructure — the infrastructure used for fine-tuning should be secure with considerations for secure data transfer, storage, and supply chain visibility.

• Transparency and Reproducibility — for security audits and model verification, it is essential to maintain precise records of fine-tuning processes, including data sources and training logs.

Using Javelin In Model Fine-Tuning Workflows

You can use Javelin in your model fine-tuning workflows. Just enable archiving on your Javelin routes, and all model interactions are automatically enabled:

‍

‍

By enabling all data collection in a central location, you drastically reduce your Security footprint and reduce data sprawl. This is highly important from a governance and responsible model usage standpoint. You can also use Javelin Archives for compliance and audits!

RAG In Action:

• Question Encoder: When a question or query is input into the RAG system, it is first passed through a question encoder to get its vector representation.

• Document Retrieval: The encoded question vector is then used to retrieve relevant document embeddings from the embedding space using a similarity measure (like cosine similarity).

• Combine and Generate: The retrieved documents (or their embeddings) are combined with the original query and fed into the generator model, producing the final response.

Prevent Sensitive Data Leaks — it is easy to provide security guardrails that inspect, redact, mask, or anonymize sensitive data before it ends in vector space. It is easy to restrict applications that generate embeddings versus applications that encode questions. You create two routes, one for the application that encodes questions and another for the application that generates embeddings. By passing requests through Javelin, you can set policies and automatic guardrails that prevent PII and PHI passthrough: