Implementing Simple Role-Based Access Control in LLM Applications: A 30-Day LLM Transformation Guide

I’m starting a new series — 30 Day LLM Transformation: Practical Insights for AI Implementation.

Welcome to the inaugural post of our series — “30 Day LLM Transformation: Practical Insights for AI Implementation.” As CEO of Mitra AI, I’ve had the privilege of training and advising numerous companies and architects on AI best practices. Through this series, I aim to share ideas and code snippets to tackle critical challenges in AI integration within enterprises.

This series is designed for those with foundational AI knowledge. I encourage you to share the challenges you’re facing in the comments, and I’ll strive to highlight solutions the industry is working on.

The Challenge of Role-Based Access Control in LLM Applications

Imagine you’re the CIO of LargeFin Corp, having developed an LLM-based Retriever Application (RAG) that transforms vast amounts of enterprise data into actionable insights.

A simple representation of a RAG query.

Through chatbots and search interfaces, your employees and customers extract immense business value. But how do you ensure that employees access only the documents they’re permitted to see?

Consider these three documents in your enterprise corpus, accessible via an internal LLM interface:

10K.PDF:
This is a public summary of the financial performance in Q1. 
Overall, the company has shown...

Q1Reflections.PPT: 
Internal analysis reveals that the Q1 performance, while positive, 
indicates potential risks...

EmergingMktStrategy.DOC: 
Confidential: The investment strategy for Q2 involves focusing on 
emerging markets...

The challenge lies in restricting access to sensitive information while allowing the intended audience to search and synthesize it. LLMs, being monolithic systems, cannot deterministically categorize data. However, utilizing the RAG pattern with a Vector DB facilitates this.

Implementing Role-Based Access Control: A Step-by-Step Guide

1. Best Practices: Start by isolating the Vector DB collection from the business logic and restricting write access.
2. Authentication Levels: Establish multiple authentication levels. In this example, we use levels 1, 2, and 3.
3. Metadata Encoding: When adding documents to the Vector DB, include the authentication level in the document’s metadata. Using ChromaDB on Langchain, we add documents as follows:

# Document 1: Public financial summary (auth_level 1)
collection.add(documents=["This is a public summary of the financial performance in Q1. Overall, the company has shown..."],
               metadatas =[{"auth_level": 1, "category": "Summary", "quarter": "Q1", "year": "2024"}],
               ids="id1")

# Document 2: Internal financial analysis (auth_level 2)
collection.add(documents=["Internal analysis reveals that the Q1 performance, while positive, indicates potential risks..."],
               metadatas =[{"auth_level": 2, "category": "Internal Analysis", "quarter": "Q1", "year": "2024"}],
               ids="id2")

# Document 3: Confidential investment strategy (auth_level 3)
collection.add(documents=["Confidential: The investment strategy for Q2 involves focusing on emerging markets..."],
               metadatas =[{"auth_level": 3, "category": "Investment Strategy", "quarter": "Q2", "year": "2024"}],
               ids="id3")

Now, provide a simple wrapper for accessing the Vector DB for your RAG:

def vector_search_with_rbac(query_text, user_auth_level):
    if not isinstance(user_auth_level, int) or user_auth_level < 0:
        raise ValueError("Invalid user_auth_level. Must be a non-negative integer.")

    try:
        results = collection.query(query_texts=[query_text],n_results=10, where={"auth_level": user_auth_level})
        return results
    except Exception as e:
        print(f"An error occurred: {e}")
        return []

This setup can be further refined with a token-based system and validation of the auth level value sent by the calling application.

Testing the System

By querying the Vector DB with access level 1, only the public summary document is returned.

Adjusting the access level yields different results, demonstrating the system’s effectiveness.

Additionally, metadata can encode department info, time, and geography for further access restrictions, offering a flexible mechanism to mirror existing roles.

Advantages of this approach:

1. Granular Access Control: By encoding access levels in the metadata of documents, the system allows for granular control over who can access what information. This is crucial for compliance with data protection regulations and internal policies.
2. Scalability: Utilizing a Vector DB to handle document access enables the system to scale efficiently as the number of documents grows, which is beneficial for large enterprises.
3. Flexibility: The approach offers flexibility in defining access levels and can be extended to include other metadata for access control, such as department, geography, and time, allowing for a highly customizable system.
4. Innovative Use of LLMs: Leveraging LLMs for retrieving and synthesizing information while controlling access through a separate layer (Vector DB) represents a smart use of current AI technologies, maintaining the strengths of LLMs in understanding and processing language.

Additionally you need to think through the following:

1. Securing your Vector DB from external intrusions and access beyond the interface above.
2. Providing a way to sync the changed access level of a file with vector db collection.

Here is a notebook for this.

Connect with me on Twitter. View me on LinkedIn.

Book a 30 min appointment for consulting.