diff --git a/notebooks/rag.ipynb b/notebooks/rag.ipynb new file mode 100644 index 0000000..c492c43 --- /dev/null +++ b/notebooks/rag.ipynb @@ -0,0 +1,507 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d597dd80", + "metadata": {}, + "source": [ + "# Fine Grained Authorization for Retrieval Augmented Generation (RAG)\n", + "\n", + "## Introduction\n", + "\n", + "Building \"enterprise-ready AI\" requires ensuring that users can only augment prompts with data they are authorized to access. Relationship-based access control (ReBAC) is particularly well-suited for fine-grained authorization in Retrieval-Augmented Generation (RAG). ReBAC allows you to generate a list of resources that a specific user has permission to access (e.g., a list of documents that a user has permission to view). This capability enables you to pre-filter vector database queries with a list of authorized object IDs, improving both efficiency and security. Moreover, ReBAC excels at fine-grained authorization because it makes decisions based on the relationships between objects, offering more precise control compared to traditional models like RBAC and ABAC.\n", + "\n", + "## RAG with ReBAC Walkthrough\n", + "\n", + "Let's walk through a practical example of implementing fine-grained authorization for RAG using Pinecone, Langchain, OpenAI, and SpiceDB.\n", + "\n", + "### Prerequisites\n", + "- Access to a [SpiceDB](https://authzed.com/spicedb) instance and API key. (This example uses insecure connections to a locally running SpiceDB instance. You can find examples of using a secure client [here](https://github.com/authzed/authzed-py?tab=readme-ov-file#basic-usage).)\n", + "- A [Pinecone](https://www.pinecone.io/) account and API key.\n", + "- An [OpenAI](https://openai.com/) account and API key.\n", + "\n", + "#### Installing Libraries\n", + "Run the following command to install the required libraries for this example:\n", + "\n", + "```pip install python-dotenv authzed pinecone pinecone[grpc] langchain_pinecone langchain langchain_openai langchain_core```\n", + "\n", + "#### Running SpiceDB\n", + "Follow the [install guide](https://authzed.com/docs/spicedb/getting-started/installing-spicedb) for your platform and run a local instance of SpiceDB.\n", + "\n", + "```spicedb serve --grpc-preshared-key rag-rebac-walkthrough```\n", + "\n", + "### Secrets\n", + "\n", + "This walkthrough requires the following environment variables. The easiest way to set these is with a .env file in your working directory:\n", + "\n", + "- ```OPENAI_API_KEY=```\n", + "- ```PINECONE_API_KEY=```\n", + "- ```SPICEDB_API_KEY=rag-rebac-walkthrough```\n", + "- ```SPICEDB_ADDR=localhost:50051```" + ] + }, + { + "cell_type": "markdown", + "id": "adb42b99", + "metadata": {}, + "source": [ + "### Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a76fa3a1", + "metadata": {}, + "outputs": [], + "source": [ + "from dotenv import load_dotenv\n", + "\n", + "#load secrets from .env file\n", + "load_dotenv()" + ] + }, + { + "cell_type": "markdown", + "id": "be225a51", + "metadata": {}, + "source": [ + "In today's scenario, we will be authorizing access to view blog articles.\n", + "\n", + "We will begin by defining the authorization logic for our example. To do this, we write a [schema](https://authzed.com/docs/spicedb/concepts/schema) to SpiceDB. The schema below defines two object types, ```user``` and ```article```. Users can relate to a document as a ```viewer``` and any user who is related to a document as a ```viewer``` can ```view``` the document." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7092cdd0", + "metadata": {}, + "outputs": [], + "source": [ + "from authzed.api.v1 import (\n", + " Client,\n", + " WriteSchemaRequest,\n", + ")\n", + "\n", + "import os\n", + "\n", + "#change to bearer_token_credentials if you are using tls\n", + "from grpcutil import insecure_bearer_token_credentials\n", + "\n", + "SCHEMA = \"\"\"definition user {}\n", + "\n", + "definition article {\n", + " relation viewer: user\n", + "\n", + " permission view = viewer\n", + "}\"\"\"\n", + "\n", + "client = Client(os.getenv('SPICEDB_ADDR'), insecure_bearer_token_credentials(os.getenv('SPICEDB_API_KEY')))\n", + "resp = client.WriteSchema(WriteSchemaRequest(schema=SCHEMA))" + ] + }, + { + "cell_type": "markdown", + "id": "2edaf49e", + "metadata": {}, + "source": [ + "Now, we write relationships to SpiceDB that specify that Tim is a viewer of document 123 and 456.\n", + "\n", + "After these relationships are written, any permission checks to SpiceDB will reflect that Tim can view documents 123 and 456." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "af300a1f", + "metadata": {}, + "outputs": [], + "source": [ + "from authzed.api.v1 import (\n", + " ObjectReference,\n", + " Relationship,\n", + " RelationshipUpdate,\n", + " SubjectReference,\n", + " WriteRelationshipsRequest,\n", + ")\n", + "\n", + "client.WriteRelationships(\n", + " WriteRelationshipsRequest(\n", + " updates=[\n", + " RelationshipUpdate(\n", + " operation=RelationshipUpdate.Operation.OPERATION_CREATE,\n", + " relationship=Relationship(\n", + " resource=ObjectReference(object_type=\"article\", object_id=\"123\"),\n", + " relation=\"viewer\",\n", + " subject=SubjectReference(\n", + " object=ObjectReference(\n", + " object_type=\"user\",\n", + " object_id=\"tim\",\n", + " )\n", + " ),\n", + " ),\n", + " ),\n", + " RelationshipUpdate(\n", + " operation=RelationshipUpdate.Operation.OPERATION_CREATE,\n", + " relationship=Relationship(\n", + " resource=ObjectReference(object_type=\"article\", object_id=\"456\"),\n", + " relation=\"viewer\",\n", + " subject=SubjectReference(\n", + " object=ObjectReference(\n", + " object_type=\"user\",\n", + " object_id=\"tim\",\n", + " )\n", + " ),\n", + " ),\n", + " ),\n", + " ]\n", + " )\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "4a8922e2", + "metadata": {}, + "source": [ + "We now define a Pinecone serverless index.\n", + "\n", + "Pinecone is a specialized database designed for handling vector-based data. Their serverless product makes it easy to get started with a vector DB." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1f3a8ba0", + "metadata": {}, + "outputs": [], + "source": [ + "from pinecone.grpc import PineconeGRPC as Pinecone\n", + "from pinecone import ServerlessSpec\n", + "import os\n", + "\n", + "pc = Pinecone(api_key=os.getenv(\"PINECONE_API_KEY\"))\n", + "\n", + "index_name = \"oscars\"\n", + "\n", + "pc.create_index(\n", + " name=index_name,\n", + " dimension=1024,\n", + " metric=\"cosine\",\n", + " spec=ServerlessSpec(\n", + " cloud=\"aws\",\n", + " region=\"us-east-1\"\n", + " )\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "b6fb61ba", + "metadata": {}, + "source": [ + "We are simulating a real-world RAG (retrieval-augmented generation) scenario by embedding a completely fictional string: \"Peyton Manning won the 2023 Oscar for best football movie.\". Since LLMs don't \"know\" this fact (as it's made up), we mimic a typical RAG case where private or unknown data augments prompts.\n", + "\n", + "In this example, we also specify metadata like article_id to track which article the string comes from. The article_id is important for linking embeddings to objects that users are authorized on." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66c1970a", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_pinecone import PineconeEmbeddings\n", + "from langchain_pinecone import PineconeVectorStore\n", + "\n", + "from langchain.schema import Document\n", + "import os\n", + "\n", + "# Create a Document object that specifies our made up article and specifies the document_id as metadata.\n", + "text = \"Peyton Manning won the 2023 Oscar for best football movie\"\n", + "metadata = {\n", + " \"article_id\": \"123\"\n", + "}\n", + "document = Document(page_content=text,metadata=metadata)\n", + "\n", + "\n", + "# Initialize a LangChain embedding object.\n", + "model_name = \"multilingual-e5-large\"\n", + "embeddings = PineconeEmbeddings(\n", + " model=model_name,\n", + " pinecone_api_key=os.environ.get(\"PINECONE_API_KEY\")\n", + ")\n", + "\n", + "namespace_name = \"oscar\"\n", + "\n", + "# Upsert the embedding into your Pinecone index.\n", + "docsearch = PineconeVectorStore.from_documents(\n", + " documents=[document],\n", + " index_name=index_name,\n", + " embedding=embeddings,\n", + " namespace=namespace_name\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "d87b8508", + "metadata": {}, + "source": [ + "### Making a request when the user is authorized to view the necessary contextual data\n", + "\n", + "Next, we will query SpiceDB for a list of documents that Tim is allowed to view. Here, we use the [LookupResources API](https://buf.build/authzed/api/docs/main:authzed.api.v1#authzed.api.v1.LookupResourcesRequest) to obtain a list of articles that ```tim``` has ```view``` permission on. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "891fca91", + "metadata": {}, + "outputs": [], + "source": [ + "from authzed.api.v1 import (\n", + " LookupResourcesRequest,\n", + " ObjectReference,\n", + " SubjectReference,\n", + ")\n", + "\n", + "subject = SubjectReference(\n", + " object=ObjectReference(\n", + " object_type=\"user\",\n", + " object_id=\"tim\",\n", + " )\n", + ")\n", + "\n", + "def lookupArticles():\n", + " return client.LookupResources(\n", + " LookupResourcesRequest(\n", + " subject=subject,\n", + " permission=\"view\",\n", + " resource_object_type=\"article\",\n", + " )\n", + " )\n", + "\n", + "resp = lookupArticles()\n", + "\n", + "authorized_articles = []\n", + "\n", + "async for response in resp:\n", + " authorized_articles.append(response.resource_object_id)\n", + "\n", + "print(\"Article IDs that Tim is authorized to view:\")\n", + "print(authorized_articles)" + ] + }, + { + "cell_type": "markdown", + "id": "11f1c9e4", + "metadata": {}, + "source": [ + "We can now issue a prompt to GPT-3.5, enhanced with relevant data that the user is authorized to access. This ensures that the response is based on information the user is permitted to view." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a28ccfae", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n", + "from langchain_pinecone import PineconeVectorStore\n", + "from langchain_core.output_parsers import StrOutputParser\n", + "from langchain.prompts import ChatPromptTemplate\n", + "from langchain.schema.output_parser import StrOutputParser\n", + "from langchain_core.runnables import (\n", + " RunnableParallel,\n", + " RunnablePassthrough\n", + ")\n", + "\n", + "# Define the ask function\n", + "def ask():\n", + " # Initialize a LangChain object for an OpenAI chat model.\n", + " llm = ChatOpenAI(\n", + " openai_api_key=os.environ.get(\"OPENAI_API_KEY\"),\n", + " model_name=\"gpt-3.5-turbo\",\n", + " temperature=0.0\n", + " )\n", + "\n", + " # Initialize a LangChain object for a Pinecone index with an OpenAI embeddings model.\n", + " knowledge = PineconeVectorStore.from_existing_index(\n", + " index_name=index_name,\n", + " namespace=namespace_name,\n", + " embedding=OpenAIEmbeddings(\n", + " openai_api_key=os.environ[\"OPENAI_API_KEY\"],\n", + " dimensions=1024,\n", + " model=\"text-embedding-3-large\"\n", + " )\n", + " )\n", + "\n", + " # Initialize a retriever with a filter that restricts the search to authorized documents.\n", + " retriever=knowledge.as_retriever(\n", + " search_kwargs={\n", + " \"filter\": {\n", + " \"article_id\":\n", + " {\"$in\": authorized_articles},\n", + " },\n", + " }\n", + " )\n", + "\n", + " # Initialize a string prompt template that let's us add context and a question.\n", + " prompt = ChatPromptTemplate.from_template(\"\"\"Answer the question below using the context:\n", + "\n", + " Context: {context}\n", + "\n", + " Question: {question}\n", + "\n", + " Answer: \"\"\")\n", + "\n", + " retrieval = RunnableParallel(\n", + " {\"context\": retriever, \"question\": RunnablePassthrough()}\n", + " )\n", + "\n", + " chain = retrieval | prompt | llm | StrOutputParser()\n", + "\n", + " question = \"\"\"Who won the 2023 Oscar for best football movie?\"\"\"\n", + "\n", + " print(\"Prompt: \\n\")\n", + " print(question)\n", + " print(chain.invoke(question))\n", + "\n", + "#invoke the ask function\n", + "ask()\n" + ] + }, + { + "cell_type": "markdown", + "id": "4900c640", + "metadata": {}, + "source": [ + "### Making a request when the user is not authorized to view the necessary contextual data\n", + "\n", + "Now, let's see what happens when Tim is not authorized to view the document.\n", + "\n", + "First, we will delete the relationship that related Tim as a viewer to document 123." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dde09648", + "metadata": {}, + "outputs": [], + "source": [ + "client.WriteRelationships(\n", + " WriteRelationshipsRequest(\n", + " updates=[\n", + " RelationshipUpdate(\n", + " operation=RelationshipUpdate.Operation.OPERATION_DELETE,\n", + " relationship=Relationship(\n", + " resource=ObjectReference(object_type=\"article\", object_id=\"123\"),\n", + " relation=\"viewer\",\n", + " subject=SubjectReference(\n", + " object=ObjectReference(\n", + " object_type=\"user\",\n", + " object_id=\"tim\",\n", + " )\n", + " ),\n", + " ),\n", + " ),\n", + " ]\n", + " )\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "31ddaef6", + "metadata": {}, + "source": [ + "Next, we will update the list of documents that Tim is authorized to view." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e744adc3", + "metadata": {}, + "outputs": [], + "source": [ + "#this function was defined above\n", + "resp = lookupArticles()\n", + "\n", + "authorized_articles = []\n", + "\n", + "async for response in resp:\n", + " authorized_articles.append(response.resource_object_id)\n", + "\n", + "print(\"Documents that Tim can view:\")\n", + "print(authorized_articles)" + ] + }, + { + "cell_type": "markdown", + "id": "6ea1c350", + "metadata": {}, + "source": [ + "Now, we can run our query again. \n", + "\n", + "Note that we no longer recieve a completion that answers our question because Tim is no longer authorized to view the document that contains the context required to answer the question." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6d5b06e4", + "metadata": {}, + "outputs": [], + "source": [ + "#this function was defined above\n", + "ask()" + ] + }, + { + "cell_type": "markdown", + "id": "403efd05", + "metadata": {}, + "source": [ + "### Clean Up\n", + "\n", + "You can now delete your Pinceone index if you'd like to." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a6e514dd", + "metadata": {}, + "outputs": [], + "source": [ + "pc.delete_index(index_name)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}