With the rise of AI-driven applications, especially those leveraging Retrieval-Augmented Generation (RAG), securing sensitive data has become paramount. Recent incidents highlight the risks associated with improper handling of vectors and embeddings. Ensuring that only authorized users can access specific documents is critical to maintaining data integrity and privacy. This becomes urgent as more companies integrate RAG into their systems, making it essential to implement robust security measures.

Understanding the RAG Application

Let’s start by examining a Ruby on Rails RAG application called Work Companion. This app acts as a chat interface for employees, providing answers based on internal company documents. The challenge here is to ensure that users can only access documents they are permitted to see, preventing any accidental exposure of sensitive information.

The RAG Process

1. Retrieval: When a user asks a question, the app searches a vector database to find the most relevant text chunks from the available documents.
2. Augmentation: These text chunks are added to the user’s original question to provide more context.
3. Generation: The combined input is sent to a Language Model (LLM) to generate a precise response.

Example Scenario

Imagine an Engineer asking about “salary bands.” In a poorly secured RAG setup, the vector search might inadvertently retrieve a snippet from a private HR document. By integrating Auth0 FGA, we ensure that only documents the user is authorized to access are considered during the retrieval process.

Setting Up the Work Companion App

Prerequisites

• Ruby 4.0.1
• PostgreSQL 17 with pgvector extension
• An OpenAI API key

Database Schema

The app uses a simple schema with three tables:

• users: Stores user information.
• documents: Contains metadata about each document.
• document_chunks: Holds the vector embeddings and uses an HNSW index for efficient searching.

Services

• RagQueryService: Manages the RAG flow.
• FgaService: Interfaces with the Auth0 FGA API to fetch user permissions.

Gems

• neighbor: Handles pgvector within ActiveRecord.
• ruby-openai: Connects to OpenAI for generating embeddings.
• openfga: Ruby SDK for interacting with Auth0 FGA.

Running the Code Sample

First, clone and install the dependencies:

git clone https://github.com/auth0-blog/ruby-rag-fga
cd ruby-rag-fga
bundle install

Next, set up PostgreSQL with the pgvector extension, create the database, and seed it with data. Follow the steps in the repo’s README. Start the server with:

rails s

Navigate to http://localhost:3000 to see the chat interface.

Adding Authentication with Auth0

Before integrating Auth0 FGA, we need to authenticate users. Auth0 handles identity management, ensuring we know exactly who the user is.

Step-by-Step Guide

1. Create an Auth0 Application:

   • Go to the Auth0 dashboard and create a new “Regular Web Application.”
   • Note down the Domain, Client ID, and Client Secret.

2. Install and Set Up the Auth0 SDK:

   • Add the following to your Gemfile:

     source "https://rubygems.org"
     # ...
     # Auth0 Authentication
     gem "omniauth-auth0", "~> 3.1"
     gem "omniauth-rails_csrf_protection", "~> 1.0"
     # ...
     

   • Install the gems:

     bundle install
     

3. Configure Auth0 in Rails:

   • Create a new initializer file config/initializers/auth0.rb:

     OmniAuth.config.logger = Rails.logger
     
     Rails.application.config.middleware.use OmniAuth::Builder do
       provider(
         :auth0,
         ENV['AUTH0_CLIENT_ID'],
         ENV['AUTH0_CLIENT_SECRET'],
         ENV['AUTH0_DOMAIN'],
         callback_path: '/auth/auth0/callback',
         failure_path: '/auth/failure'
       )
     end
     
     OmniAuth.config.on_failure = Proc.new do |env|
       OmniAuth::FailureEndpoint.new(env).redirect_to_failure
     end
     

4. Set Environment Variables:

   • Add your Auth0 credentials to your .env file:

5. Create Routes for Authentication:

   • Add the following routes to config/routes.rb:

     get '/auth/auth0/callback' => 'auth0#callback'
     get '/auth/failure' => 'auth0#failure'
     get '/logout' => 'auth0#logout'
     

6. Implement Auth0 Controller:

   • Create a controller app/controllers/auth0_controller.rb:

     class Auth0Controller < ApplicationController
       def callback
         session[:userinfo] = request.env['omniauth.auth']
         redirect_to root_path
       end
     
       def failure
         @error_type = request.params['error_type']
         @error_msg = request.params['error_description']
         flash[:alert] = "Authentication error: #{@error_msg}."
         redirect_to root_path
       end
     
       def logout
         reset_session
         redirect_to logout_url.to_s
       end
     
       private
     
       def logout_url
         domain = ENV['AUTH0_DOMAIN']
         client_id = ENV['AUTH0_CLIENT_ID']
         params = {
           returnTo: root_url,
           client_id: client_id
         }
     
         URI::HTTPS.build(host: domain, path: '/v2/logout', query: params.to_query)
       end
     end
     

7. Update Views:

   • Add login and logout links in your views. For example, in app/views/layouts/application.html.erb:

     <% if session[:userinfo] %>
       <p>Welcome <%= session[:userinfo][:info][:name] %>!</p>
       <%= link_to 'Logout', logout_path %>
     <% else %>
       <%= link_to 'Login', '/auth/auth0' %>
     <% end %>
     

Integrating Auth0 FGA for Fine-Grained Authorization

Now that users are authenticated, we need to ensure they can only access documents they are authorized to see. Auth0 FGA provides the necessary tools to implement fine-grained authorization.

Step-by-Step Guide

1. Set Up Auth0 FGA:

   • Go to the Auth0 dashboard and create a new FGA application.
   • Note down the API URL and credentials.

2. Install the OpenFGA SDK:

   • Add the following to your Gemfile:

     gem 'openfga', '~> 0.1'
     

   • Install the gem:

     bundle install
     

3. Configure OpenFGA in Rails:

   • Create a new initializer file config/initializers/openfga.rb:

     OpenFGA.configure do |config|
       config.api_url = ENV['OPENFGA_API_URL']
       config.client_id = ENV['OPENFGA_CLIENT_ID']
       config.client_secret = ENV['OPENFGA_CLIENT_SECRET']
     end
     

4. Set Environment Variables:

   • Add your OpenFGA credentials to your .env file:

5. Implement FgaService:

   • Create a service app/services/fga_service.rb:

     class FgaService
       def self.get_allowed_documents(user_id)
         client = OpenFGA::Client.new
         response = client.read(
           type: 'document',
           relation: 'viewer',
           user: "user:#{user_id}"
         )
         response.objects.map { |obj| obj.split(':').last }
       end
     end
     

6. Update RagQueryService:

   • Modify app/services/rag_query_service.rb to filter documents based on user permissions:

     class RagQueryService
       def initialize(user_id)
         @user_id = user_id
       end
     
       def query(question)
         allowed_document_ids = FgaService.get_allowed_documents(@user_id)
         chunks = DocumentChunk.where(document_id: allowed_document_ids)
         embeddings = chunks.pluck(:vector)
         # Perform vector search and augmentation
         # Send to LLM for generation
       end
     end
     

7. Integrate with Controllers:

   • Ensure the user ID is passed to RagQueryService. For example, in app/controllers/chats_controller.rb:

     class ChatsController < ApplicationController
       before_action :authenticate_user!
     
       def create
         user_id = session[:userinfo][:uid]
         service = RagQueryService.new(user_id)
         response = service.query(params[:question])
         render json: { response: response }
       end
     
       private
     
       def authenticate_user!
         redirect_to '/auth/auth0' unless session[:userinfo]
       end
     end
     

Key Takeaways

• Authentication: Use Auth0 to manage user identities securely.
• Authorization: Implement Auth0 FGA to enforce fine-grained access control.
• Data Integrity: Ensure that only authorized users can access specific documents, preventing data leakage.

Conclusion

Securing Ruby on Rails RAG applications is crucial to protect sensitive data. By integrating Auth0 for authentication and Auth0 FGA for fine-grained authorization, you can ensure that only authorized users access specific documents. This setup not only enhances security but also improves the user experience by providing accurate and relevant information.