RAG authorization system in LangGraph using Cerbos and Pinecone

Welcome to my second article on the security of RAG-based AI agents. In my previous guide, we explored how to build a permission system using LangChain and Chroma DB. This time, we’ll dive into implementing an authorization system for your RAG app using a new set of tools: LangGraph, Pinecone, and Cerbos. Along the way, you’ll learn how to:

• Set up a Pinecone dataset for RAG.
• Build an Enterprise Resource Planning RAG application with LangGraph.
• Implement Role-Based Access Control and Attribute-Based Access Control using Cerbos.
• Practical examples of integrating all tools together to have a secure Agentic RAG application (LangGraph, Pinecone, and Cerbos).

Why should you build AI agents with security in mind? Since 2022, I’ve worked on RAG apps like chatbots, virtual assistants, and enterprise AI solutions, and I’ve seen firsthand how much my clients cared about protecting their sensitive data. Agentic RAG applications are powerful because they use real business data, but that also means they handle sensitive information stored in vector databases. That is why access control (strong RAG authorization) is essential—it ensures only authorized users can access, add, or modify data, keeping everything secure.

So, let’s dive in and unlock the potential of secure and efficient RAG applications!

Prerequisites

We'll start with the basic set up. Following are the prerequisites to run scripts in this guide:

• Sign-up for a Pinecone account and retrieve your Pinecone API key

• Install the Langchain, LangGraph and Cerbos libraries to access Pinecone database, the OpenAI API, and the Cerbos server.

!pip install -qU langchain-pinecone pinecone-notebooks
!pip install -qU langchain-openai
!pip install -qU langchain-core
!pip install -qU langchain-community langgraph
!pip install -qU cerbos

The next is to import the relevant classes and functions from the above modules:

import getpass
import os
import time
import sqlite3
import pandas as pd
import time
from uuid import uuid4
from typing_extensions import List, TypedDict
from google.protobuf.struct_pb2 import Value
from IPython.display import Image, display

from pinecone import Pinecone, ServerlessSpec

from langchain_openai import OpenAIEmbeddings
from langchain_openai import ChatOpenAI
from langchain_pinecone import PineconeVectorStore
from langchain_core.documents import Document
from langchain import hub
from langgraph.graph import START, END, StateGraph
from langchain import hub

from cerbos.sdk.grpc.client import CerbosClient
from cerbos.engine.v1.engine_pb2 import Principal, Resource

Setting Up Pinecone Dataset for RAG

To keep this guide practical, we’ll build a RAG-based AI agent for Enterprise Resource Planning. It’s a great example because sales, partner, employee, and company resource data are sensitive. You want to safeguard this data when users ask the AI agent questions like “What information do you store about our business partners?” and “What do we record in the sales order?”

To make it all work, we’ll use a sample sales dataset from SAP’s bike store. The first step is to set up a Pinecone index and store the data.

Configure and Create a Pinecone Index

An index in Pinecone is a data structure that stores vectors. You can perform various functions on an index, such as inserting and deleting items, performing a vector similarity search to retrieve similar vectors, adding vectors with metadata, and so on. You can create a Pinecone index by logging into your Pinecone account. However, for this section, I will create an index using the Pinecone Python API.

The following script creates a Pinecone class object. You must pass your Pinecone API key while creating a Pinecone object.

if not os.getenv("PINECONE_API_KEY"):
    os.environ["PINECONE_API_KEY"] = getpass.getpass("Enter your Pinecone API key: ")

pinecone_api_key = os.environ.get("PINECONE_API_KEY")

pc = Pinecone(api_key=pinecone_api_key)

Next, we will configure and create a Pinecone index. The script below creates a Pinecone index named erp-data-index. The index will be able to store 3072 dimensional vectors. We use these dimensions since the embedding model that we will use creates 3072 dimensional vectors. We will use the cosine metric for similarity search. Finally, we define the cloud server and region for our Pinecone index.

index_name = "erp-data-index"  
existing_indexes = [index_info["name"] for index_info in pc.list_indexes()]

if index_name not in existing_indexes:
    pc.create_index(
        name=index_name,
        dimension=3072,
        metric="cosine",
        spec=ServerlessSpec(cloud="aws", region="us-east-1"),
    )
    while not pc.describe_index(index_name).status["ready"]:
        time.sleep(1)

index = pc.Index(index_name)

We will also create a Pinecone vector store that LangGraph will use to connect to the Pinecone index. The following script below defines the embedding model and the vector store. We will use the text-embedding-3-large model from OpenAI. I find this to be the best embedding model accross various applications.

if not os.environ.get("OPENAI_API_KEY"):
  os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")

embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

vector_store = PineconeVectorStore(index=index, embedding=embeddings)

Inserting Data in a Pinecone Index

So now we will insert the Bike Sales ERP dataset in the Pinecone index we created. I decided to use sales data for bikes but it can work with any other type of sales dataset. You can download our bikes dataset from Kaggle.

Let's see how the dataset looks. The following script fetches all the CSV files from the dataset folder and prints their names and shapes.

## dataset download link: https://www.kaggle.com/datasets/yasinnaal/bikes-sales-sample-data?resource=download

folder_path = "erp_data"

csv_files = [file for file in os.listdir(folder_path) if file.endswith('.csv')]

datasets = {}

for file in csv_files:
    file_path = os.path.join(folder_path, file)
    try:
        df = pd.read_csv(file_path, encoding='latin1')
        file_name_without_extension = os.path.splitext(file)[0]
        datasets[file_name_without_extension] = df
    except Exception as e:
        print(f"Error reading {file}: {e}")

for file_name, df in datasets.items():
    print(f"File Name: {file_name}")
    print(f"DataFrame Shape: {df.shape}")
    print("-" * 40)

Output:

Let's see data in SalesOrder and Employees table:

datasets["SalesOrders"].head(10)

Output:

datasets["Employees"].head(10)

Output:

We will store each row in a table as a vector in the Pinecone index. To do so, we will first convert each row in a table into a LangChain document, where the page_content attribute of the document will store all the row values. We will append the column name before each row value and concatenate the values into a single string. In addition, we will add the table name as the source in the document's metadata.

The following script iterates through all the rows in our dataset's CSV files and converts them into LangChain documents.

documents = []

for file_name, df in datasets.items():
    for _, row in df.iterrows():
        row_data = ", ".join([f"{col}: {row[col]}" for col in df.columns])

        document = Document(
            page_content=row_data,
            metadata={"source": file_name}
        )

        documents.append(document)
        
print(len(documents))
print(documents[0])

Output:

The above output shows that we have 2476 documents. You can also partially see how the first document looks.

The next step is to add these values to the vector store, which will, in turn, store the values in the Pinecone index (erp-data-index) that we attached to this vector store.

uuids = [str(uuid4()) for _ in range(len(documents))]

vector_store.add_documents(documents=documents, ids=uuids)

If you open your Pinecone dashboard, you should see the erp-data-index as in the following screenshot:

You can see a sample document containing the metadata, text, and vector values.

Let's search a few documents in the index. To do so, you can use the similarity_search() function. Here we retrieve the top 5 most similar documents using K=5. You can use any other number but 5 is the most commonly used number for RAG operations.

results = vector_store.similarity_search(
    "What information do you store about product text?",
    k=5,
    filter= {"source":"ProductTexts"}
)
for res in results:
    print(f"* {res.page_content} [{res.metadata}]")

Output:

The above script retrieves the documents with the source attribute ProductTexts.

You can also search documents in an index by converting a vectore store into a retriever. This approach is more useful for RAG operations, as seen in a later section.

retriever = vector_store.as_retriever()

retriever.invoke("What information do you store about product text?",
                filter= {"source":"ProductTexts"})

Output:

RAG with LangGraph on Pinecone Dataset

Now that we have created a Pinecone index and a LangChain vector store, we are ready to create a LangGraph agent for retrieval augmented generation. So, let's dive in.

Though you can use LangChain to create RAG apps, I recommend LangGrpah, which is the state-of-the-art, more agentic approach to creating chatbot applications integrating RAG.

Defining Graph State and Prompts

You must define the graph state before you create a LangGraph agent. The graph's state contains information shared among all the nodes and edges of a graph.

The following script defines the state of our RAG agent.

# Define state for application
class State(TypedDict):
    question: str
    context: List[Document]
    answer: str

The state consists of a question that will contain the user input, the context, which we will retrieve from the Pinecone index, and the answer, which will be generated via a large language model.

Next, we will define a RAG prompt and the LLM (GPT-4o) that we will use to generate a response. We will use a built-in RAG prompt from the LangChain hub.

llm = ChatOpenAI(model="gpt-4o")

example_messages = prompt.invoke(
    {"context": "(context goes here)", "question": "(question goes here)"}
).to_messages()

assert len(example_messages) == 1
print(example_messages[0].content)

Output:

You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
Question: (question goes here) 
Context: (context goes here) 
Answer:

The above output shows the RAG prompt we will use to generate responses from the LLM.

Retriever and Response Generation Nodes

Next, we will define retriever and response generation nodes for our LangGraph.

The retriever node will extract question from the graph's state and pass it to the retrieve.invoke() function to retrieve similar documents from the Pinecone index. The documents are stored in the graph's state's context attribute.

def retrieve(state: State):
    retrieved_docs = retriever.invoke(state["question"])
    return {"context": retrieved_docs}

Let's test the retrieve() function.

input = {"question":"What information do you store about our business partners?"}
retrieve(input)

Output:

You can see that the retrieve() function returns the documents relevant to the user query.

Next, we will define the generate() function, which extracts question and context from the graph's state, embeds them in the prompt, and passes them to the LLM to generate a response. The response is stored in the graph's state's answer attribute.

This might seem to take longer to set up than getting physical authorization with paperwork and long queues. But the on the bright side, you only have to do it only once ;).

def generate(state: State):
    docs_content = "\n\n".join(doc.page_content for doc in state["context"])
    messages = prompt.invoke({"question": state["question"], "context": docs_content})
    response = llm.invoke(messages)
    return {"answer": response.content}

Putting it Together: Creating a LangGraph

Now that we have retriever and generation nodes, we will integrate them in a sequence to create a LangGraph agent.

graph_builder = StateGraph(State).add_sequence([retrieve, generate])
graph_builder.add_edge(START, "retrieve")
graph = graph_builder.compile()

display(Image(graph.get_graph().draw_mermaid_png()))

Output:

Our graph starts from the retrieve node and finishes on the generate node. We will extract the final response from the graph's answer attribute.

Let's test the graph. Imagine a company's manager using our RAG system wants to retrieve detailed information about an employee named Javas. He can execute the following query.

input = {"question":"I want all the information about the employee Javas"}
result = graph.invoke(input)

print(f'Context: {result["context"]}\n\n')
print(f'Answer: {result["answer"]}')

Output:

The output shows the context documents and the final response. The response is spot on; you can verify it by looking at the Employees.csv file.

We only need the final answer to generate a response. Let's ask another question:

input = {"question":"What information do you have about our business partners, can you give names and details of some of them"}
result = graph.invoke(input)

print(f'Answer: {result["answer"]}')

Output:

Answer: Some of the business partners include Avante Grande Bikes, Amaze Bikes Inc, eBike 100, and A to Z Fitness. Avante Grande Bikes is a Ltd. company based in the UK. Amaze Bikes Inc and eBike 100 are both Inc. companies based in the US and Australia, respectively, while A to Z Fitness is based in the United Arab Emirates.

Again, RAG has been spot on.

We have successfully developed a RAG agent capable of generating responses based on information in a Pinecone index. But is this enough? What if an unauthorized person gets access to your RAG application and retrieves information about all of your business partners, employees, and the products you sell? That can be disastrous, and you certainly do not want that. This is where RAG authorization comes in, and our demo app is a great way to show why permission management matters

Adding Authorization to LangGraph RAG with Cerbos

Cerbos is an open-source, language-agnostic authorization layer designed to streamline the implementation of authorization layers. Cerbos offers enhanced security, scalability, and simplified management of access control policies.

Cerbos uses a declarative approach to define authorization policies, separating the authorization logic from your application's core functionality. It is highly scalable, ensuring efficient handling of high-volume authorization requests. Having personally used Cerbos in my RAG applications, I can confidently affirm that it is among the most efficient and user-friendly access control tools available for RAG today.

So in this section, you'll see how to use Cerbos for implementing RBAC and ABAC within the LangGraph RAG app introduced earlier in the first section.

Setting up the Environment

Using the Cerbos authorization layer in your application requires installing the Cerbos server, as explained in the official documentation. Since we developed our RAG-based AI application in LangGraph, you must install the Cerbos Python SDK to integrate Cerbos authorization in our application.

You can use pip to install the SDK as follows:

pip install cerbos

Finally, we will import the required Cerbos modules into our application in addition to the libraries you imported in the previous sections.

from cerbos.sdk.grpc.client import CerbosClient
from cerbos.engine.v1.engine_pb2 import Principal, Resource

Example of RBAC on LangGraph RAG with Cerbos

RBAC lets users with specific roles perform certain actions on restricted resources. To set up RBAC with Cerbos, you first need to define authorization policies in a .yaml file. For this example, I’ll keep the policy structure simple, but in real-world AI agents (especially for enterprise use cases) the roles and policies can get much more complex. This article can be a nice starting point for mapping business requirements to authorization policies!

So, the following defines Cerbos authorization policy for the resource name vector_store. It defines two actions: all and retrieve_employee. The principals (users in this case) with roles admin and CEO will be allowed to perform all actions, whereas the hr_manager role will be allowed the retrieve_employe action on the vector_store resource.

apiVersion: "api.cerbos.dev/v1"
resourcePolicy:
  resource: "vector_store"
  version: "default"
  rules:

    - actions: ["all"]
      effect: EFFECT_ALLOW
      roles: ["admin", "CEO"]

    - actions: ["retrieve_employee"]
      effect: EFFECT_ALLOW
      roles: ["hr_manager"]
      

Next, while starting the Cerbos server you must specify the path to your yaml file containing policies. For example, if your policies are in cerbos-quickstart/policies/resource.document.yaml file, you will start your Cerbos server with the following command:

docker run --name cerbos -d -v $(pwd)/cerbos-quickstart/policies:/policies -p 3592:3592 -p 3593:3593 ghcr.io/cerbos/cerbos:0.39.0

Next, to integrate the authorization logic in the RAG application, we will add the principal, resource, is_authorized, and meta_filters attributes to the graph state.

In addition, we will define the authorize() method that will be added as the starting node in our graph. Inside the authorize method we create an object of the CerbosClient class.

Next, we pass the action names, and the principal and resource objects from the state graph to the is_allowed() method. The is_allowed() method returns true if a principal can perform the specified action on the resource.

In the authorize() method, we check authorization for the all and retrieve_employee actions. If a principal is authorized to perform all actions, the meta_filter of the state graph will be none since the principal can access all documents in the Pinecone index. For the retrieve_employee action, we set a filter to only return the Employee documents from the index since we want the retrieve_employee action to only access the Emoloyees documents. This is just an example, you can set other filters as per your application logic.

If a principal can access a resource the is_authorized attribute of the state graph is set to True. Finally, the authorize() method returns the the values of the question, meta_filters, and is_authorized attributes of the state graph.

If a principal cannot access the resource, the answer attribute of the state graph is assigned a message that explains that the user cannot access the resource. In addition the value of the is_authorized attribute is set to False. Nice and simple 🙌.

The following script shows the updated state graph and implementation for the authrize() method.

# Define state for application
class State(TypedDict):
    question: str
    context: List[Document]
    answer: str
    principal: Principal # adding Cerbos principal to graph state
    resource: Resource # adding Cerbos resource to graph state
    is_authorized: bool
    meta_filters: List[str]


def authorize(state: State):

        with CerbosClient("localhost:3592", tls_verify=False) as client:

            # the action is allowed for users with admin and CEO roles only
            if client.is_allowed("all", state["principal"], state["resource"]):
                filters = None
                return {"question": state["question"], "is_authorized": True, "meta_filters":filters}

            # this action is allowed for users with roles "hr_manager" and it can only retrieve data from "Employees" table
            elif client.is_allowed("retrieve_employee",state["principal"], state["resource"]):
                filters= ["Employees"]
                return {"question": state["question"], "is_authorized": True, "meta_filters":filters}

            else:
                unauthorized_message = "You are not authorized to access the resource; hence no information can be retrieved."
                return {"answer": unauthorized_message, "is_authorized": False}

In the retrieve() method we return context documents based on the values in the meta_filters. There will be no change in the generate() method.

Next, we add the authorize, retrieve, and generate nodes to our graph, and set the authorize node as the starting point.

We also add a conditional edge from the authorize node which based on the value of the is_authorize attribute of the state graph, goes directly to the end of graph, or to the retrieve node.

def retrieve(state: State):
    if state["meta_filters"] is None:
        retrieved_docs = retriever.invoke(state["question"])
    else:
        retrieved_docs = retriever.invoke(state["question"],
                                         filter = {"source": {"$in":state["meta_filters"]}}
                                                    )
    return {"context": retrieved_docs}

def generate(state: State):
    docs_content = "\n\n".join(doc.page_content for doc in state["context"])
    messages = prompt.invoke({"question": state["question"], "context": docs_content})
    response = llm.invoke(messages)
    return {"answer": response.content}

# Create the graph builder
graph_builder = StateGraph(State)

# Add nodes
graph_builder.add_node("authorize", authorize)
graph_builder.add_node("retrieve", retrieve)
graph_builder.add_node("generate", generate)

# Add edges
graph_builder.add_edge(START, "authorize")
graph_builder.add_conditional_edges(
    "authorize",
    lambda x: x["is_authorized"],
    {
        True: "retrieve",
        False: END
    }
)
graph_builder.add_edge("retrieve", "generate")
graph_builder.add_edge("generate", END)

# Compile the graph
graph = graph_builder.compile()

# Display the graph
display(Image(graph.get_graph().draw_mermaid_png()))

Output:

The above output shows our updated graph. If the is_authorized attribute is set to False, the graph goes directly to the end node. Otherwise, it retrieves context related to the user query, generates a response, and reaches the end node.

Let's test our graph now. We will create three users (principal) with roles admin, CEO and hr_manager, and a vector_store type resource.

# Define Principals with different roles
principal_user1 = engine_pb2.Principal(
    id="user1",
    roles=["admin"],
    policy_version="default",
)

principal_user2 = engine_pb2.Principal(
    id="user2",
    roles=["CEO"],
    policy_version="default",
)

principal_user3 = engine_pb2.Principal(
    id="user3",
    roles=["hr_manager"],
    policy_version="default",
)

resource_rbac = engine_pb2.Resource(
    id="erp_vector_store",
    kind="vector_store",
)

First we will see what documents can the user1 with role admin can retrieve. Pass the principal, the resource and the question to our graph.

input = {"question":"What information do you have about our business partners, can you give names and details including websites of some of them?",
        "principal": principal_user1,
        "resource":resource_rbac}

result = graph.invoke(input)
print(f'Answer: {result["answer"]}')

Output:

Answer: Some of your business partners include Amaze Bikes Inc, All For Bikes, and eBike 100. Amaze Bikes Inc can be found at [amazebikes.com](http://www.amazebikes.com), All For Bikes at [all4bikes.com](http://www.all4bikes.com), and eBike 100 at [ebike100.com](http://www.ebike100.com).

The above output shows that user1 was able to retrieve an answer to a question about business partners since it has access to all the documents in the index, including documents from the BusinessPartners file.

Let's ask the same question using the user3 principal with hr_manager role.

input = {"question":"What information do you have about our business partners, can you give names and details of some of them?",
        "principal": principal_user3,
        "resource":resource_rbac}

result = graph.invoke(input)
print(f'Answer: {result["answer"]}')

Output:

Answer: I don't have any information about business partners or their details in the provided context.

The above output shows that the user with hr_manager role could not retrieve answers for questions related to business partners. So, our AI agent safeguarded our sensitive data, great job!

Next, we will ask a question about employees using the hr_manager role.

input = {"question":"I need detailed information about Javas?",
        "principal": principal_user3,
        "resource":resource_rbac}

result = graph.invoke(input)
print(f'Answer: {result["answer"]}')

Output:

Answer: Javas Hegde is an employee with ID 5, whose contact information includes the phone number 02224135120 and email address javas.hegde@itelo.info. His login name is javasm, and his employee record is valid from January 1, 2000, to December 31, 9999.

Since user3 has the hr_mnager role, he can access the Employee documents in the Pinecone index, as you can see from the above output. Again, everything worked as expected.

Let's see an example of attribute-based access control with Cerbos.

Example of ABAC on LangGraph RAG with Cerbos

ABAC allows principals to perform actions on resources based on their attributes. For example, only a principal whose department matches the resource's department can access the resource.

Let's modify our policy document to add actions that are allowed based on resources' attribute values.

We will have three actions in total: all, retrieve_employee, and retrieve_products.

The all action will allow users with admin and CEO roles to perform any action on the vector_store type resource.

The retrieve_employee action will authorize users with hr_manager roles to access the vector_store resource whose department attribute is set to Employees.

Similarly, the retrieve_products action will authorize the product_manager role to access the vector_store resource for the departments:Products, ProductTexts, ProductCategories, ProductCategoryText department.

apiVersion: "api.cerbos.dev/v1"
resourcePolicy:
  resource: "vector_store"
  version: "default"
  rules:
    - actions: ["all"]
      effect: EFFECT_ALLOW
      roles: ["admin", "CEO"]

    - actions: ["retrieve_employee"]
      effect: EFFECT_ALLOW
      roles: ["hr_manager"]
      condition:
        match:
          expr: request.resource.attr.department == "Employees"

    - actions: ["retrieve_products"]
      effect: EFFECT_ALLOW
      roles: ["product_manager"]
      condition:
        match:
          expr: request.resource.attr.department in ["Products", "ProductTexts", "ProductCategories", "ProductCategoryText"]

Let's test the above policy. We will create three principals and three resources to test the three resource scenarios in the policy.

Notice, that in this case we define resources using the PlanResourcesInput.Resource class since we will use the Cerbos Plan Resource API to implement authorization on RAG queries.

The plan resource API is particularly useful when you want to authorize principals to access a resource based on resource's attributes.

principal_ceo = engine_pb2.Principal(
    id="company ceo",
    roles=["CEO"],
    policy_version="default"
)

principal_hr_manager = engine_pb2.Principal(
    id="user_hr_manager",
    roles=["hr_manager"],
    policy_version="default"
)

principal_product_manager = engine_pb2.Principal(
    id="product_manager",
    roles=["product_manager"],
    policy_version="default"
)


resource_employee_department = engine_pb2.PlanResourcesInput.Resource(
    kind="vector_store",
    attr={"department": Value(string_value="Employees")}
)

resource_product_department = engine_pb2.PlanResourcesInput.Resource(
    kind="vector_store",
    attr={"department": Value(string_value="Products")}
)

resource_all_departments = engine_pb2.PlanResourcesInput.Resource(
    kind="vector_store"
)

To make the code more dynamic, we will add an action attribute in the graphs's state which corresponds to the action performed by the principal.

class State(TypedDict):
    question: str
    context: List[Document]
    answer: str
    action: str
    principal: Principal # adding Cerbos principal to graph state
    resource: Resource # adding Cerbos resource to graph state
    is_authorized: bool
    meta_filters: List[str]

We will also modify the authorize() function. In this case, we will use the client.plan_resources() class to create a query plan.

Next, we will check if the plan.filter.kind attribute equals KIND_ALWAYS_ALLOWED, we authorize the principal to access a resource. Else if the filter kind is KIND_ALWAYS_DENIED, we deny the access to the resource.

In addition, if a principal is authorized acces to a resource, we extract the value of its department attribute and pass it to the RAG meta_filters. This allows to restrict the RAG output to the relevant department only.

If a resource does not have a value for the department attribute, we set the filters to None.

def authorize(state: State):
    
    with CerbosClient("localhost:3592", tls_verify=False) as client:

        plan = client.plan_resources(action =  state["action"], 
                                     principal = state["principal"], 
                                     resource = state["resource"])
                          
        if plan.filter.kind == plan.filter.KIND_ALWAYS_ALLOWED:

            department_attr = state["resource"].attr.get("department")
            if department_attr.string_value != "":
                filters = [department_attr.string_value]
            else:
                filters = None
                
            return {"question": state["question"], "is_authorized": True, "meta_filters":filters}
        
        elif plan.filter.kind == plan.filter.KIND_ALWAYS_DENIED:
            unauthorized_message = "You are not authorized to access the resource; hence no information can be retrieved."
            return {"answer": unauthorized_message, "is_authorized": False}
                                     

The rest of the code remains the same for the graph.

# Create the graph builder
graph_builder = StateGraph(State)

# Add nodes
graph_builder.add_node("authorize", authorize)
graph_builder.add_node("retrieve", retrieve)
graph_builder.add_node("generate", generate)

# Add edges
graph_builder.add_edge(START, "authorize")
graph_builder.add_conditional_edges(
    "authorize",
    lambda x: x["is_authorized"],
    {
        True: "retrieve",
        False: END
    }
)
graph_builder.add_edge("retrieve", "generate")
graph_builder.add_edge("generate", END)

# Compile the graph
graph = graph_builder.compile()

Let's test the ABAC policies we defined earlier. Let's ask a question using principal_ceo on resource resource_all_departments.

input = {"question":"Who is Javas. Provide me all details about him.",
         "action": "all",
        "principal": principal_ceo,
        "resource": resource_all_departments}

result = graph.invoke(input)
print(f'Answer: {result["answer"]}')

Output:

Answer: Javas Hegde is an employee with the ID 5. His contact details include a phone number 02224135120 and an email address javas.hegde@itelo.info. His login name is javasm, and his employee record is valid from January 1, 2000, until December 31, 9999.

Since the principal_ceo can perform all actions on the resource_all_departments resource, you can see a response to a question related to an employee.

Let's ask a different question that requires accessing information from a different set of documents.

input = {"question":"What are some of our business partners?.",
         "action": "all",
        "principal": principal_ceo,
        "resource": resource_all_departments}

result = graph.invoke(input)
print(f'Answer: {result["answer"]}')

Output:

Answer: Some of our business partners include Amaze Bikes Inc, Avante Grande Bikes, Atlanta Corp Inc, and Tires On Fire.

You can see that theprincipal_ceo can successfully retrieve information about business partners as well.

Now let's try accessing the resource_all_departments using the principal_product_manager.

input = {"question":"What are some of our business partners?.",
         "action": "retrieve_products",
        "principal": principal_product_manager,
        "resource": resource_all_departments}

result = graph.invoke(input)
print(f'Answer: {result["answer"]}')

Output:

Answer: You are not authorized to access the resource; hence no information can be retrieved.

You can see that the princpal_product_manager cannot access the resource_all_departments since it can only perform retrieve_products action on the resources from departments Products, ProductTexts, ProductCategories, or ProductCategoryText.

On the other hand, the principal_product_manager will be able to access the resource_product_department from the Products department.

input = {"question":"What information do we store about our products?",
         "action": "retrieve_products",
        "principal": principal_product_manager,
        "resource": resource_product_department}

result = graph.invoke(input)
print(f'Answer: {result["answer"]}')

Output:

Answer: We store information such as PRODUCTID, TYPECODE, PRODCATEGORYID, CREATEDBY, CREATEDAT, CHANGEDBY, CHANGEDAT, SUPPLIER_PARTNERID, TAXTARIFFCODE, QUANTITYUNIT, WEIGHTMEASURE, WEIGHTUNIT, CURRENCY, PRICE, and potentially some dimensional data like WIDTH, DEPTH, HEIGHT, and PRODUCTPICURL, although the latter are not available (nan) in the given data.

Let's see one more example. Let's try to access the resource_employee_department using the principal_hr_manager.

input = {"question":"What information do we store about our products?",
         "action": "retrieve_employee",
        "principal": principal_hr_manager,
        "resource": resource_employee_department}

result = graph.invoke(input)
print(f'Answer: {result["answer"]}')

Output:

Answer: I don't know. The provided context only contains information about employees, not products.

You can see that the principal_hr_manager was able to access the resource_employee_department but since this resource only contains information about employees, the query about products was not answered. Congrats, your AI agent is secure!