--- title: "How to implement an authorization system for RAG using LangGraph, Cerbos and Pinecone" description: "Dive into implementing an authorization system for your RAG app using 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." author: "Usman Malik" date: "2025-02-20T18:10:00.000Z" canonical: "https://www.cerbos.dev/blog/rag-authorization-system-langgraph-cerbos-pinecone" image: "https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/RAG_Authorization_System_in_Lang_Graph_using_Cerbos_and_Pinecone_803e40c171.png" tags: ["engineering","guide"] source: "https://www.cerbos.dev/blog/rag-authorization-system-langgraph-cerbos-pinecone" --- # How to implement an authorization system for RAG using LangGraph, Cerbos and Pinecone Welcome to my second article on the security of RAG-based AI agents. In [my previous guide](https://www.cerbos.dev/blog/authorization-for-rag-applications-langchain-chromadb-cerbos), 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](https://www.cerbos.dev/)). 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. ```python= !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: ```python= 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. ```python= 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. ```python= 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. ```python= 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](https://www.kaggle.com/datasets/yasinnaal/bikes-sales-sample-data?resource=download). 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. ```python= ## 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:** ![Inserting Data in a Pinecone Index.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/Inserting_Data_in_a_Pinecone_Index_6b36c8056b.png) Let's see data in `SalesOrder` and `Employees` table: ```python= datasets["SalesOrders"].head(10) ``` **Output:** ![data in `SalesOrder` and `Employees` table.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/data_in_Sales_Order_and_Employees_table_0c2d6de5d3.png) ```python= datasets["Employees"].head(10) ``` **Output:** ![employees.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/employees_bfdf34208e.png) 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. ```python= 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:** ![documents.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/documents_17539e0096.png) 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. ```python= 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: ![erp-data-index.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/erp_data_index_c2b781f313.png) 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. ```python= 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:** ![similarity search.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/similarity_search_af98776d54.png) 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. ```python= retriever = vector_store.as_retriever() retriever.invoke("What information do you store about product text?", filter= {"source":"ProductTexts"}) ``` **Output:** ![producttexts.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/producttexts_080cd98401.png) ## 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. ```python= # 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](https://smith.langchain.com/hub/rlm/rag-prompt). ```python= 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. ```python= def retrieve(state: State): retrieved_docs = retriever.invoke(state["question"]) return {"context": retrieved_docs} ``` Let's test the `retrieve()` function. ```python= input = {"question":"What information do you store about our business partners?"} retrieve(input) ``` **Output:** ![retrieve.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/retrieve_d5eea7187c.png) 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 ;). ```python= 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. ```python= 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:** ![langgraph.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/langgraph_bfcf0ff143.png) 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. ```python= 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:** ![graph test.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/graph_test_7cafc057da.png) 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: ```python= 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](https://docs.cerbos.dev/cerbos/latest/installation/container). Since we developed our RAG-based AI application in LangGraph, you must install the [Cerbos Python SDK](https://github.com/cerbos/cerbos-sdk-python) 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. ```python= 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](https://www.cerbos.dev/blog/mapping-business-requirements-to-authorization-policy) 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. ```python= # 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. ```python= 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:** ![updated graph.png](https://stylish-appliance-1c1cc1c30d.media.strapiapp.com/updated_graph_fafa84254b.png) 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. ```python= # 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. ```python= 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. ```python= 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. ```python= 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](https://docs.cerbos.dev/cerbos/latest/api/?_gl=1*15axkdm*_up*MQ..*_ga*MTYwNDc4ODYwNS4xNzM4MjI1NTg0*_ga_8G3G3MS838*MTczODIyNTU4My4xLjAuMTczODIyNTU4My4wLjAuMA..#resources-query-plan) 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. ```python= 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`. ```python= 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. ```python= # 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`. ```python= 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. ```python= 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 the`principal_ceo` can successfully retrieve information about business partners as well. Now let's try accessing the `resource_all_departments` using the `principal_product_manager`. ```python= 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. ```python= 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`. ```python 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!