{ "cells": [ { "cell_type": "markdown", "id": "XwR-PYCFu0Nd", "metadata": { "id": "XwR-PYCFu0Nd" }, "source": [ "# Building a Role-Based RAG Pipeline with Redis\n", "\n", "This notebook demonstrates a simplified setup for a **Role-Based Retrieval Augmented Generation (RAG)** pipeline, where:\n", "\n", "1. Each **User** has one or more **roles**.\n", "2. Knowledge base **Documents** in Redis are tagged with the official roles that can access them (`allowed_roles`).\n", "3. A unified **query flow** ensures a user only sees documents that match at least one of their roles.\n", "\n", "![Role Based RAG](https://raw.githubusercontent.com/redis-developer/redis-ai-resources/main/assets/role-based-rag.png)" ] }, { "cell_type": "markdown", "id": "58823e66", "metadata": { "id": "58823e66" }, "source": [ "\n", "## Let's Begin!\n", "\"Open" ] }, { "cell_type": "code", "execution_count": 7, "id": "4e0aa177", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "4e0aa177", "outputId": "0ba61596-b3e4-442f-cd9c-8b480f1c52d1" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Note: you may need to restart the kernel to use updated packages.\n" ] } ], "source": [ "%pip install -q \"redisvl>=0.6.0\" openai langchain-community pypdf langchain-text-splitters" ] }, { "cell_type": "markdown", "id": "fXsGCsLQu0Ne", "metadata": { "id": "fXsGCsLQu0Ne" }, "source": [ "## 1. High-Level Data Flow & Setup\n", "\n", "1. **User Creation & Role Management**\n", " - A user is stored at `user:{user_id}` in Redis with a JSON structure containing the user’s roles.\n", " - We can create, update, or delete users as needed.\n", " - **This serves as a simple look up layer and should NOT replace your production-ready auth API flow**\n", "\n", "2. **Document Storage**\n", " - Documents chunks are stored at `doc:{doc_id}:{chunk_id}` in Redis as JSON.\n", " - Each document chunk includes fields such as `doc_id`, `chunk_id`, `content`, `allowed_roles`, and an `embedding` (for vector similarity).\n", "\n", "3. **Querying / Search**\n", " - User roles are retrieved from Redis.\n", " - We perform a vector similarity search (or any other type of retrieval) on the documents.\n", " - We filter the results so that only documents whose `allowed_roles` intersect with the user’s roles are returned.\n", "\n", "4. **RAG Integration**\n", " - The returned documents can be fed into a Large Language Model (LLM) to provide context and generate an answer.\n", "\n", "First, we’ll set up our Python environment and Redis connection.\n" ] }, { "cell_type": "markdown", "id": "73c33af6", "metadata": { "id": "73c33af6" }, "source": [ "### Download Documents\n", "Running remotely or in collab? Run this cell to download the necessary datasets." ] }, { "cell_type": "code", "execution_count": null, "id": "48971c52", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "48971c52", "outputId": "e17d146a-43be-41fb-b029-f330d79f1a65" }, "outputs": [], "source": [ "# NBVAL_SKIP\n", "!git clone https://github.com/redis-developer/redis-ai-resources.git temp_repo\n", "!mkdir -p resources\n", "!mv temp_repo/python-recipes/RAG/resources/aapl-10k-2023.pdf resources/\n", "!mv temp_repo/python-recipes/RAG/resources/2022-chevy-colorado-ebrochure.pdf resources/\n", "!rm -rf temp_repo" ] }, { "cell_type": "markdown", "id": "993371a2", "metadata": { "id": "993371a2" }, "source": [ "### Run Redis Stack\n", "\n", "For this tutorial you will need a running instance of Redis if you don't already have one.\n", "\n", "#### For Colab\n", "Use the shell script below to download, extract, and install [Redis Stack](https://redis.io/docs/getting-started/install-stack/) directly from the Redis package archive." ] }, { "cell_type": "code", "execution_count": 4, "id": "8edc5862", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "8edc5862", "outputId": "df2643ed-2422-4ee5-bd42-bec17b405eec" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "deb [signed-by=/usr/share/keyrings/redis-archive-keyring.gpg] https://packages.redis.io/deb jammy main\n", "Starting redis-stack-server, database path /var/lib/redis-stack\n" ] } ], "source": [ "# NBVAL_SKIP\n", "%%sh\n", "curl -fsSL https://packages.redis.io/gpg | sudo gpg --dearmor -o /usr/share/keyrings/redis-archive-keyring.gpg\n", "echo \"deb [signed-by=/usr/share/keyrings/redis-archive-keyring.gpg] https://packages.redis.io/deb $(lsb_release -cs) main\" | sudo tee /etc/apt/sources.list.d/redis.list\n", "sudo apt-get update > /dev/null 2>&1\n", "sudo apt-get install redis-stack-server > /dev/null 2>&1\n", "redis-stack-server --daemonize yes" ] }, { "cell_type": "markdown", "id": "bc571319", "metadata": { "id": "bc571319" }, "source": [ "#### For Alternative Environments\n", "There are many ways to get the necessary redis-stack instance running\n", "1. On cloud, deploy a [FREE instance of Redis in the cloud](https://redis.com/try-free/). Or, if you have your\n", "own version of Redis Enterprise running, that works too!\n", "2. Per OS, [see the docs](https://redis.io/docs/latest/operate/oss_and_stack/install/install-stack/)\n", "3. With docker: `docker run -d --name redis-stack-server -p 6379:6379 redis/redis-stack-server:latest`" ] }, { "cell_type": "code", "execution_count": 2, "id": "qU49fNVnu0Nf", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "qU49fNVnu0Nf", "outputId": "4d2f34c3-6179-4f1d-eff7-5e8e9d8fd58b" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Successfully connected to Redis\n" ] } ], "source": [ "import os\n", "\n", "from redis import Redis\n", "\n", "# Replace values below with your own if using Redis Cloud instance\n", "REDIS_HOST = os.getenv(\"REDIS_HOST\", \"localhost\") # ex: \"redis-18374.c253.us-central1-1.gce.cloud.redislabs.com\"\n", "REDIS_PORT = os.getenv(\"REDIS_PORT\", \"6379\") # ex: 18374\n", "REDIS_PASSWORD = os.getenv(\"REDIS_PASSWORD\", \"\") # ex: \"1TNxTEdYRDgIDKM2gDfasupCADXXXX\"\n", "\n", "# If SSL is enabled on the endpoint, use rediss:// as the URL prefix\n", "REDIS_URL = f\"redis://:{REDIS_PASSWORD}@{REDIS_HOST}:{REDIS_PORT}\"\n", "\n", "# Connect to Redis (adjust host/port if needed)\n", "redis_client = Redis.from_url(REDIS_URL)\n", "redis_client.ping()\n", "\n", "print(\"Successfully connected to Redis\")" ] }, { "cell_type": "code", "execution_count": 10, "id": "4be07006", "metadata": {}, "outputs": [], "source": [ "import getpass\n", "\n", "if \"OPENAI_API_KEY\" not in os.environ:\n", " os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OPENAI_API_KEY :\")" ] }, { "cell_type": "markdown", "id": "aqzMteQsu0Nf", "metadata": { "id": "aqzMteQsu0Nf" }, "source": [ "## 2. User Management\n", "\n", "Below is a simple `User` class that stores a user in Redis as JSON. We:\n", "\n", "- Use a Redis key of the form `user:{user_id}`.\n", "- Store fields like `user_id`, `roles`, etc.\n", "- Provide CRUD methods (Create, Read, Update, Delete) for user objects.\n", "\n", "**Data Structure Example**\n", "```json\n", "{\n", " \"user_id\": \"alice\",\n", " \"roles\": [\"finance\", \"manager\"]\n", "}\n", "```\n", "\n", "We'll also include some basic checks to ensure we don't add duplicate roles, handle empty role lists, etc.\n" ] }, { "cell_type": "code", "execution_count": 3, "id": "38pdjXJvu0Nf", "metadata": { "id": "38pdjXJvu0Nf" }, "outputs": [], "source": [ "from typing import List, Optional\n", "from enum import Enum\n", "\n", "\n", "class UserRoles(str, Enum):\n", " FINANCE = \"finance\"\n", " MANAGER = \"manager\"\n", " EXECUTIVE = \"executive\"\n", " HR = \"hr\"\n", " SALES = \"sales\"\n", " PRODUCT = \"product\"\n", "\n", "\n", "class User:\n", " \"\"\"\n", " User class for storing user data in Redis.\n", "\n", " Each user has:\n", " - user_id (string)\n", " - roles (list of UserRoles)\n", "\n", " Key in Redis: user:{user_id}\n", " \"\"\"\n", " def __init__(\n", " self,\n", " redis_client: Redis,\n", " user_id: str,\n", " roles: Optional[List[UserRoles]] = None\n", " ):\n", " self.redis_client = redis_client\n", " self.user_id = user_id\n", " self.roles = roles or []\n", "\n", " @property\n", " def key(self) -> str:\n", " return f\"user:{self.user_id}\"\n", "\n", " def exists(self) -> bool:\n", " \"\"\"Check if the user key exists in Redis.\"\"\"\n", " return self.redis_client.exists(self.key) == 1\n", "\n", " def create(self):\n", " \"\"\"\n", " Create a new user in Redis. Fails if user already exists.\n", " \"\"\"\n", " if self.exists():\n", " raise ValueError(f\"User {self.user_id} already exists.\")\n", "\n", " self.save()\n", "\n", " def save(self):\n", " \"\"\"\n", " Save (create or update) the user data in Redis.\n", " If user does not exist, it will be created.\n", " \"\"\"\n", " data = {\n", " \"user_id\": self.user_id,\n", " \"roles\": [UserRoles(role).value for role in set(self.roles)] # ensure roles are unique and convert to strings\n", " }\n", " self.redis_client.json().set(self.key, \".\", data)\n", "\n", " @classmethod\n", " def get(cls, redis_client: Redis, user_id):\n", " \"\"\"\n", " Retrieve a user from Redis.\n", " \"\"\"\n", " key = f\"user:{user_id}\"\n", " data = redis_client.json().get(key)\n", " if not data:\n", " return None\n", " # Convert string roles back to UserRoles enum\n", " roles = [UserRoles(role) for role in data.get(\"roles\", [])]\n", " return cls(redis_client, data[\"user_id\"], roles)\n", "\n", " def update_roles(self, roles: List[UserRoles]):\n", " \"\"\"\n", " Overwrite the user's roles in Redis.\n", " \"\"\"\n", " self.roles = roles\n", " self.save()\n", "\n", " def add_role(self, role: UserRoles):\n", " \"\"\"Add a single role to the user.\"\"\"\n", " if role not in self.roles:\n", " self.roles.append(role)\n", " self.save()\n", "\n", " def remove_role(self, role: UserRoles):\n", " \"\"\"Remove a single role from the user.\"\"\"\n", " if role in self.roles:\n", " self.roles.remove(role)\n", " self.save()\n", "\n", " def delete(self):\n", " \"\"\"Delete this user from Redis.\"\"\"\n", " self.redis_client.delete(self.key)\n", "\n", " def __repr__(self):\n", " return f\"\"\n" ] }, { "cell_type": "markdown", "id": "FNQxAaoCxPN7", "metadata": { "id": "FNQxAaoCxPN7" }, "source": [ "### Example usage of User class" ] }, { "cell_type": "code", "execution_count": 4, "id": "_WcOlgVyu0Ng", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "_WcOlgVyu0Ng", "outputId": "0776fa25-513b-445b-d46d-35d9333b3a75" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "User 'alice' created.\n", "Retrieved: \n", "After adding 'executive': \n", "After removing 'manager': \n" ] } ], "source": [ "# Example usage of the User class\n", "\n", "# Let's create a new user\n", "alice = User(redis_client, \"alice\", roles=[\"finance\", \"manager\"])\n", "\n", "# We'll save the user in Redis\n", "try:\n", " alice.create()\n", " print(\"User 'alice' created.\")\n", "except ValueError as e:\n", " print(e)\n", "\n", "# Retrieve the user\n", "alice_obj = User.get(redis_client, \"alice\")\n", "print(\"Retrieved:\", alice_obj)\n", "\n", "# Add another role\n", "alice_obj.add_role(\"executive\")\n", "print(\"After adding 'executive':\", alice_obj)\n", "\n", "# Remove a role\n", "alice_obj.remove_role(\"manager\")\n", "print(\"After removing 'manager':\", alice_obj)\n" ] }, { "cell_type": "code", "execution_count": 5, "id": "c911e892", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "c911e892", "outputId": "df4666ff-97ce-4e75-d70c-75fe5d9e6703" }, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Take a peek at the user object itself\n", "alice" ] }, { "cell_type": "code", "execution_count": 6, "id": "P3j6yu8l87j3", "metadata": { "id": "P3j6yu8l87j3" }, "outputs": [], "source": [ "# Create one more user\n", "larry = User(redis_client, \"larry\", roles=[\"product\"])\n", "larry.create()" ] }, { "cell_type": "markdown", "id": "Y7B4l7XVx5md", "metadata": { "id": "Y7B4l7XVx5md" }, "source": [ ">💡 Using a cloud DB? Take a peek at your instance using [RedisInsight](https://redis.io/insight) to see what user data is in place." ] }, { "cell_type": "markdown", "id": "aCXYFXu0u0Ng", "metadata": { "id": "aCXYFXu0u0Ng" }, "source": [ "## 3. Document Management (Using LangChain)\n", "\n", "Here, we'll use **LangChain** for document loading, chunking, and vectorizing. Then, we’ll **store documents** in Redis as JSON. Each document will look like:\n", "\n", "```json\n", "{\n", " \"doc_id\": \"123\",\n", " \"chunk_id\": \"123\",\n", " \"path\": \"resources/doc.pdf\",\n", " \"title\": \"Quarterly Finance Report\",\n", " \"content\": \"Some text...\",\n", " \"allowed_roles\": [\"finance\", \"executive\"],\n", " \"embedding\": [0.12, 0.98, ...] \n", "}\n", "```" ] }, { "cell_type": "markdown", "id": "d3cJ5DSP5vXt", "metadata": { "id": "d3cJ5DSP5vXt" }, "source": [ "### Building a document knowledge base\n", "We will create a `KnowledgeBase` class to encapsulate document processing logic and search. The class will handle:\n", "1. Document ingest and chunking\n", "2. Role tagging with a simple str-based rule (likely custom depending on use case)\n", "3. Retrieval over the entire document corpus adhering to provided user roles\n" ] }, { "cell_type": "code", "execution_count": 11, "id": "67d38524", "metadata": { "id": "67d38524" }, "outputs": [], "source": [ "from typing import List, Optional, Dict, Any, Set\n", "from pathlib import Path\n", "import uuid\n", "\n", "from langchain_community.document_loaders import PyPDFLoader\n", "from langchain_text_splitters import RecursiveCharacterTextSplitter\n", "from redisvl.index import SearchIndex\n", "from redisvl.query import VectorQuery\n", "from redisvl.query.filter import FilterExpression, Tag\n", "from redisvl.utils.vectorize import OpenAITextVectorizer\n", "\n", "\n", "class KnowledgeBase:\n", " \"\"\"Manages document processing, embedding, and storage in Redis.\"\"\"\n", "\n", " def __init__(\n", " self,\n", " redis_client,\n", " embeddings_model: str = \"text-embedding-3-small\",\n", " chunk_size: int = 2500,\n", " chunk_overlap: int = 100\n", " ):\n", " self.redis_client = redis_client\n", " self.embeddings = OpenAITextVectorizer(model=embeddings_model)\n", " self.text_splitter = RecursiveCharacterTextSplitter(\n", " chunk_size=chunk_size,\n", " chunk_overlap=chunk_overlap,\n", " )\n", "\n", " # Initialize document search index\n", " self.index = self._create_search_index()\n", "\n", " def _create_search_index(self) -> SearchIndex:\n", " \"\"\"Create the Redis search index for documents.\"\"\"\n", " schema = {\n", " \"index\": {\n", " \"name\": \"docs\",\n", " \"prefix\": \"doc\",\n", " \"storage_type\": \"json\"\n", " },\n", " \"fields\": [\n", " {\n", " \"name\": \"doc_id\",\n", " \"type\": \"tag\",\n", " },\n", " {\n", " \"name\": \"chunk_id\",\n", " \"type\": \"tag\",\n", " },\n", " {\n", " \"name\": \"allowed_roles\",\n", " \"path\": \"$.allowed_roles[*]\",\n", " \"type\": \"tag\",\n", " },\n", " {\n", " \"name\": \"content\",\n", " \"type\": \"text\",\n", " },\n", " {\n", " \"name\": \"embedding\",\n", " \"type\": \"vector\",\n", " \"attrs\": {\n", " \"dims\": self.embeddings.dims,\n", " \"distance_metric\": \"cosine\",\n", " \"algorithm\": \"flat\",\n", " \"datatype\": \"float32\"\n", " }\n", " }\n", " ]\n", " }\n", " index = SearchIndex.from_dict(schema, redis_client=self.redis_client)\n", " index.create()\n", " return index\n", "\n", " def ingest(self, doc_path: str, allowed_roles: Optional[List[str]] = None) -> str:\n", " \"\"\"\n", " Load a document, chunk it, create embeddings, and store in Redis.\n", " Returns the document ID.\n", " \"\"\"\n", " # Generate document ID\n", " doc_id = str(uuid.uuid4())\n", " path = Path(doc_path)\n", "\n", " if not path.exists():\n", " raise FileNotFoundError(f\"Document not found: {doc_path}\")\n", "\n", " # Load and chunk document\n", " loader = PyPDFLoader(str(path))\n", " pages = loader.load()\n", " chunks = self.text_splitter.split_documents(pages)\n", " print(f\"Extracted {len(chunks)} for doc {doc_id} from file {str(path)}\", flush=True)\n", "\n", " # If roles not provided, determine from filename\n", " if allowed_roles is None:\n", " allowed_roles = self._determine_roles(path)\n", "\n", " # Prepare chunks for Redis\n", " data, keys = [], []\n", " for i, chunk in enumerate(chunks):\n", " # Create embedding w/ openai\n", " embedding = self.embeddings.embed(chunk.page_content)\n", "\n", " # Prepare chunk payload\n", " chunk_id = f\"chunk_{i}\"\n", " key = f\"doc:{doc_id}:{chunk_id}\"\n", " data.append({\n", " \"doc_id\": doc_id,\n", " \"chunk_id\": chunk_id,\n", " \"path\": str(path),\n", " \"content\": chunk.page_content,\n", " \"allowed_roles\": list(allowed_roles),\n", " \"embedding\": embedding,\n", " })\n", " keys.append(key)\n", "\n", " # Store in Redis\n", " _ = self.index.load(data=data, keys=keys)\n", " print(f\"Loaded {len(chunks)} chunks for document {doc_id}\")\n", " return doc_id\n", "\n", " def _determine_roles(self, file_path: Path) -> Set[str]:\n", " \"\"\"Determine allowed roles based on file path and name patterns.\"\"\"\n", " # Customize based on use case and business logic\n", " ROLE_PATTERNS = {\n", " ('10k', 'financial', 'earnings', 'revenue'):\n", " {'finance', 'executive'},\n", " ('brochure', 'spec', 'product', 'manual'):\n", " {'product', 'sales'},\n", " ('hr', 'handbook', 'policy', 'employee'):\n", " {'hr', 'manager'},\n", " ('sales', 'pricing', 'customer'):\n", " {'sales', 'manager'}\n", " }\n", "\n", " filename = file_path.name.lower()\n", " roles = {\n", " role for terms, roles in ROLE_PATTERNS.items()\n", " for role in roles\n", " if any(term in filename for term in terms)\n", " }\n", " return roles or {'executive'}\n", "\n", " @staticmethod\n", " def role_filter(user_roles: List[str]) -> FilterExpression:\n", " \"\"\"Generate a Redis filter based on provided user roles.\"\"\"\n", " return Tag(\"allowed_roles\") == user_roles\n", "\n", " def search(self, query: str, user_roles: List[str], top_k: int = 5) -> List[Dict[str, Any]]:\n", " \"\"\"\n", " Search for documents matching the query and user roles.\n", " Returns list of matching documents.\n", " \"\"\"\n", " # Create query vector\n", " query_vector = self.embeddings.embed(query)\n", "\n", " # Build role filter\n", " roles_filter = self.role_filter(user_roles)\n", "\n", " # Execute search\n", " return self.index.query(\n", " VectorQuery(\n", " vector=query_vector,\n", " vector_field_name=\"embedding\",\n", " filter_expression=roles_filter,\n", " return_fields=[\"doc_id\", \"chunk_id\", \"allowed_roles\", \"content\"],\n", " num_results=top_k,\n", " dialect=4\n", " )\n", " )\n" ] }, { "cell_type": "markdown", "id": "YsBuAa_q9QU_", "metadata": { "id": "YsBuAa_q9QU_" }, "source": [ "Load a document into the knowledge base." ] }, { "cell_type": "code", "execution_count": 12, "id": "s1LDdWhKu0Nh", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "s1LDdWhKu0Nh", "outputId": "66e1105e-78ba-425a-8156-c810c7c9054a" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Extracted 34 for doc c72e57bf-3eb2-478f-bfdc-09da0c7c7133 from file resources/2022-chevy-colorado-ebrochure.pdf\n", "Loaded 34 chunks for document c72e57bf-3eb2-478f-bfdc-09da0c7c7133\n", "Loaded all chunks for c72e57bf-3eb2-478f-bfdc-09da0c7c7133\n" ] } ], "source": [ "kb = KnowledgeBase(redis_client)\n", "\n", "doc_id = kb.ingest(\"resources/2022-chevy-colorado-ebrochure.pdf\")\n", "print(f\"Loaded all chunks for {doc_id}\", flush=True)" ] }, { "cell_type": "markdown", "id": "-Ekqkf1fu0Nh", "metadata": { "id": "-Ekqkf1fu0Nh" }, "source": [ "## 4. User Query Flow\n", "\n", "Now that we have our User DB and our Vector DB loaded in Redis. We will perform:\n", "\n", "1. **Vector Similarity Search** on `embedding`.\n", "2. A metadata **Filter** based on `allowed_roles`.\n", "3. Return top-k matching document chunks.\n", "\n", "This is implemented below.\n" ] }, { "cell_type": "code", "execution_count": 13, "id": "WpvrXmluu0Nh", "metadata": { "id": "WpvrXmluu0Nh" }, "outputs": [], "source": [ "def user_query(user_id: str, query: str):\n", " \"\"\"\n", " Placeholder for a search function.\n", " 1. Load the user's roles.\n", " 2. Perform a vector search for docs.\n", " 3. Filter docs that match at least one of the user's roles.\n", " 4. Return top-K results.\n", " \"\"\"\n", " # 1. Load & validate user roles\n", " user_obj = User.get(redis_client, user_id)\n", " if not user_obj:\n", " raise ValueError(f\"User {user_id} not found.\")\n", "\n", " roles = set([role.value for role in user_obj.roles])\n", " if not roles:\n", " raise ValueError(f\"User {user_id} does not have any roles.\")\n", "\n", " # 2. Retrieve document chunks\n", " results = kb.search(query, roles)\n", "\n", " if not results:\n", " raise ValueError(f\"No available documents found for {user_id}\")\n", "\n", " return results" ] }, { "cell_type": "markdown", "id": "qQS1BLwGBVDA", "metadata": { "id": "qQS1BLwGBVDA" }, "source": [ "### Search examples\n", "\n", "Search with a non-existent user." ] }, { "cell_type": "code", "execution_count": null, "id": "wYishsNy6lty", "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 287 }, "id": "wYishsNy6lty", "outputId": "dfa5a8b5-d926-4e94-e8a1-ecceb51ccff5" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "User tyler not found.\n" ] } ], "source": [ "# NBVAL_SKIP\n", "try:\n", " results = user_query(\"tyler\", query=\"What is the make and model of the vehicle here?\")\n", "except ValueError as e:\n", " # If a non-existent user is used for search, a ValueError should be raised \n", " print(e)\n" ] }, { "cell_type": "markdown", "id": "0af59693", "metadata": {}, "source": [ "Create user for Tyler." ] }, { "cell_type": "code", "execution_count": 18, "id": "ZNgxlQSvChx7", "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 329 }, "id": "ZNgxlQSvChx7", "outputId": "d59aad34-2d24-4c87-dd42-b9a44ccaf26b" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "'engineering' is not a valid UserRoles\n" ] } ], "source": [ "# NBVAL_SKIP\n", "try:\n", " tyler = User(redis_client, \"tyler\", roles=[\"sales\", \"engineering\"])\n", " tyler.create()\n", "except ValueError as e:\n", " # ValueError should be raised here as \"engineering\" was not defined in the UserRoles Enum\n", " print(e)" ] }, { "cell_type": "code", "execution_count": 19, "id": "WWVJF0UVCt4d", "metadata": { "collapsed": true, "id": "WWVJF0UVCt4d" }, "outputs": [], "source": [ "# Try again but this time with valid roles\n", "tyler = User(redis_client, \"tyler\", roles=[\"sales\"])\n", "tyler.create()" ] }, { "cell_type": "code", "execution_count": 20, "id": "DXEyktWLC1cC", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "DXEyktWLC1cC", "outputId": "dbb6e93f-3b81-4c14-f329-daf97a613c89" }, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tyler" ] }, { "cell_type": "code", "execution_count": 21, "id": "O0K_rdC7C6OH", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "O0K_rdC7C6OH", "outputId": "f823f253-cf42-4975-f711-6391b36f83bd" }, "outputs": [ { "data": { "text/plain": [ "[{'id': 'doc:c72e57bf-3eb2-478f-bfdc-09da0c7c7133:chunk_13',\n", " 'vector_distance': '0.606250166893',\n", " 'doc_id': '[\"c72e57bf-3eb2-478f-bfdc-09da0c7c7133\"]',\n", " 'chunk_id': '[\"chunk_13\"]',\n", " 'allowed_roles': '[\"sales\",\"product\"]',\n", " 'content': '[\"LT FEATURES\\\\nJet Black Leather-Appointed Front Seating Surfaces 4 (shown)Jet Black ClothAvailable on Extended Cab, Crew Cab Short Box and Crew Cab Long Box. In addition to standard equipment, \\\\nselect LT features include:\\\\nMECHANICAL\\\\nEngine – 2.5L DOHC 4-cylinder with Variable Valve Timing (VVT) and Direct Injection\\\\n– 3.6L DOHC V6 with Variable Valve Timing (VVT) and Direct Injection (Crew Cab 4x4 and Crew Cab Long Box 2WD models)\\\\nTransfer case – 2-speed, electronic Autotrac® with rotary controls; includes Neutral position for dinghy towing (4x4 models)\\xa0\\\\nTransmission – 6-speed automatic, electronically controlled with overdrive\\\\n– 8-speed automatic, electronically controlled with overdrive, Tow/Haul mode and Hitch Guidance 1 (Crew Cab 4x4 and Crew Cab \\\\nLong Box 2WD models)\\\\nEXTERIOR\\\\nMirrors – power-adjustable, manual-folding with body-color caps; includes driver spotter mirror\\\\nRecovery hooks 2 – front (4x4 models)\\\\nTires – 255/65R17 all-season\\\\nINTERIOR\\\\nAir conditioning – manual, single-zone\\\\nCruise control – electronic with Set and Resume Speed\\\\nEntertainment – Chevrolet Infotainment 3 system 3 with 8-inch diagonal color touch-screen\\\\nMirror – rearview, auto-dimming\\\\nRear Vision Camera 1\\\\nRemote Keyless Entry – extended-range\\\\nSeat – driver 6-way power adjuster\\xa0\\\\nSteering column – manual tilt and telescopic\\\\nSteering wheel – leather-wrapped with mounted audio, cruise and phone controls\\\\n1 Safety or driver assistance features are no substitute for the driver’s responsibility to operate the vehicle in a safe manner. Read the vehicle Owner’s Manual for important feature limitations and information. 2 To avoid the risk of injury, never use recovery \\\\nhooks to tow a vehicle. For more information, see the Recovery Hooks section of your Owner’s Manual. 3 Chevrolet Infotainment System functionality varies by model. Full functionality requires compatible Bluetooth and smartphone, and USB connectivity for \\\\nsome devices. 4 Available on Crew Cab models only. Requires available Luxury Package.\\\\n17\\\\\" Blade Silver Metallic-Painted Aluminum Wheels \\\\nStandard on LT\\\\n18\\\\\" Dark Argent Metallic-Painted Aluminum Wheels \\\\nAvailable on LT\\\\n18\\\\\" Black-Painted Aluminum Wheels with Red Accents\\\\nAvailable on LT with Redline Edition\\\\nPreproduction models shown. Actual production models may vary. Some features shown may have limited, late, or no availability. See dealer for feature availability.\"]'},\n", " {'id': 'doc:c72e57bf-3eb2-478f-bfdc-09da0c7c7133:chunk_11',\n", " 'vector_distance': '0.613473713398',\n", " 'doc_id': '[\"c72e57bf-3eb2-478f-bfdc-09da0c7c7133\"]',\n", " 'chunk_id': '[\"chunk_11\"]',\n", " 'allowed_roles': '[\"sales\",\"product\"]',\n", " 'content': '[\"Dark Ash ClothDark Ash Vinyl (shown)\\\\nWORK TRUCK (WT) FEATURES\\\\n17\\\\\" Ultra Silver Metallic-Painted Steel Wheels \\\\nStandard on WT\\\\n18\\\\\" Dark Argent Metallic-Painted Aluminum Wheels\\\\nAvailable on WT with Custom Special Edition\\\\nAvailable on Extended Cab, Crew Cab Short Box and Crew Cab Long Box. In addition to standard equipment, \\\\nselect WT features include:\\\\nMECHANICAL\\\\nEngine – 2.5L DOHC 4-cylinder with Variable Valve Timing (VVT) and Direct Injection\\\\n– 3.6L DOHC V6 with Variable Valve Timing (VVT) and Direct Injection (Crew Cab 4x4 and Crew Cab Long Box 2WD models)\\\\nTransfer case – 2-speed, electronic with rotary controls; includes Neutral position for dinghy towing (4x4 models)\\xa0\\\\nTransmission – 6-speed automatic, electronically controlled with overdrive\\\\n– 8-speed automatic, electronically controlled with overdrive, Tow/Haul mode and Hitch Guidance 1 (Crew Cab 4x4 and Crew Cab \\\\nLong Box 2WD models)\\\\nEXTERIOR\\\\nMirrors – manual-adjustable, manual-folding with Black caps\\\\nRecovery hooks 2 – front (4x4 models)\\\\nTires – 255/65R17 all-season\\\\nINTERIOR\\\\nAir conditioning – manual, single-zone\\\\nEntertainment – Chevrolet Infotainment 3 system 3 with 7-inch diagonal color touch-screen\\\\nRear Vision Camera 1\\\\nSeat – driver 4-way power adjuster with manual recline\\\\nSteering column – manual tilt\\xa0\\\\n1 Safety or driver assistance features are no substitute for the driver’s responsibility to operate the vehicle in a safe manner. Read the vehicle Owner’s Manual for important feature limitations and information. 2 To avoid the risk of injury, never use recovery \\\\nhooks to tow a vehicle. For more information, see the Recovery Hooks section of your Owner’s Manual. 3 Chevrolet Infotainment System functionality varies by model. Full functionality requires compatible Bluetooth and smartphone, and USB connectivity for \\\\nsome devices.\\\\nPreproduction models shown. Actual production models may vary. Some features shown may have limited, late, or no availability. See dealer for feature availability.\"]'},\n", " {'id': 'doc:c72e57bf-3eb2-478f-bfdc-09da0c7c7133:chunk_19',\n", " 'vector_distance': '0.624065339565',\n", " 'doc_id': '[\"c72e57bf-3eb2-478f-bfdc-09da0c7c7133\"]',\n", " 'chunk_id': '[\"chunk_19\"]',\n", " 'allowed_roles': '[\"sales\",\"product\"]',\n", " 'content': '[\"models only; LT requires available Safety Package, LT Convenience Package and Trailering Package; \\\\nLT 2WD also requires Short Box; Z71 requires available 4x4, Safety Package and Trailering Package)\\\\n●\\\\n●\\\\n—\\\\n●\\\\n●\\\\n●\\\\n—\\\\n●\\\\n●\\\\n—\\\\n●\\\\n●\\\\nFrame – fully boxed ● ● ● ●\\\\nHill Descent Control — — ● ●\\\\nStabiliTrak – Electronic Stability Control System and Traction Control; includes Hill Start Assist ● ● ● ●\\\\nSteering – rack-and-pinion with Electric Power Steering (EPS) assist ● ● ● ●\\\\nSuspension – Z71 Off-Road Package\\\\n– Multimatic DSSV™ Damping System\\\\n—\\\\n—\\\\n—\\\\n—\\\\n●\\\\n—\\\\n—\\\\n●\\\\nTow/Haul mode – includes Hitch Guidance 2 with dynamic trailering assist guideline (standard on Crew Cab \\\\nLong Box WT and LT; available with 3.6L V6 engine on Extended Cab and Crew Cab Short Box WT and LT; \\\\navailable with Duramax 2.8L Turbo-Diesel I-4 engine on Extended Cab and Crew Cab Short Box LT) ● ● ● ●\\\\nTrailer brake controller – integrated\\xa0(included with available Duramax 2.8L Turbo-Diesel I-4 engine; \\\\nrequires available Trailering Package with 3.6L V6 engine) ● ● ● ●\\\\nTransfer case – electric, 2-speed (4x4 models)\\\\n– electric, 2-speed, Autotrac (4x4 models)\\\\n– transfer case shield (4x4 models)\\\\n●\\\\n—\\\\n—\\\\n—\\\\n●\\\\n—\\\\n—\\\\n●\\\\n●\\\\n—\\\\n●\\\\n●\\\\nTransmission – 6-speed automatic, electronically controlled with overdrive (included with 2.5L DOHC \\\\n4-cylinder engine and available Duramax 2.8L Turbo-Diesel I-4 engine)\\\\n– 8-speed automatic, electronically controlled with overdrive (included with 3.6L V6 engine)\\\\n●\\\\n●\\\\n●\\\\n●\\\\n●\\\\n●\\\\n●\\\\n●\\\\n1 A NOTE ON CHILD SAFETY: Always use seat belts and the correct child restraint for your child’s age and size, even with airbags. Even in vehicles equipped with the Passenger Sensing System, children are safer when properly secured in a rear seat in the \\\\nappropriate infant, child or booster seat. Never place a rear-facing infant restraint in the front seat of any vehicle equipped with an active frontal airbag. See your vehicle Owner’s Manual and the child safety seat instructions for more safety information. \\\\n2 Safety or driver assistance features are no substitute for the driver’s responsibility to operate the vehicle in a safe manner. Read the vehicle Owner’s Manual for important feature limitations and information. 3 Chevrolet Infotainment System functionality \\\\nvaries by model. Full functionality requires compatible Bluetooth and smartphone, and USB connectivity for some devices. 4 Map coverage available in the United States, Puerto Rico and Canada. 5 When you select a monthly plan within 30 days of activating\"]'}]" ] }, "execution_count": 21, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Query with valid user\n", "results = user_query(\n", " tyler.user_id,\n", " query=\"What is the make and model of the vehicle here?\"\n", ")\n", "results[:3]" ] }, { "cell_type": "markdown", "id": "454ce79b", "metadata": {}, "source": [ "Search with a valid user, but incorrect roles." ] }, { "cell_type": "code", "execution_count": 22, "id": "irqwMseYDSS_", "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 394 }, "id": "irqwMseYDSS_", "outputId": "acb3fe4b-c451-464f-c214-8a90d835f9ef" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " \n", "\n", "No available documents found for alice\n" ] } ], "source": [ "# NBVAL_SKIP\n", "print(alice, \"\\n\")\n", "\n", "# Query with valid user\n", "try:\n", " results = user_query(\n", " alice.user_id, query=\"What is the make and model of the vehicle here?\"\n", " )\n", " print(results)\n", "except ValueError as e:\n", " # If no documents exist for the user within the knowledge base, a ValueError should be raised\n", " print(e)" ] }, { "cell_type": "markdown", "id": "c309b53d", "metadata": { "id": "c309b53d" }, "source": [ "Empty results because there are no documents available for Alice to view. Add some." ] }, { "cell_type": "code", "execution_count": 23, "id": "0e5e990b", "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 1000 }, "id": "0e5e990b", "outputId": "b0b1bc64-6b01-47d3-feb4-3d6d1cc8e38d" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Extracted 155 for doc 5b403537-786d-4ba4-b7e6-5d67383dd3fd from file resources/aapl-10k-2023.pdf\n", "Loaded 155 chunks for document 5b403537-786d-4ba4-b7e6-5d67383dd3fd\n" ] }, { "data": { "text/plain": [ "'5b403537-786d-4ba4-b7e6-5d67383dd3fd'" ] }, "execution_count": 23, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Add a document that Alice will have access to\n", "kb.ingest(\"resources/aapl-10k-2023.pdf\")" ] }, { "cell_type": "code", "execution_count": 24, "id": "9fcf8cc0", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "9fcf8cc0", "outputId": "bce13955-7d37-472b-f820-5588cd3986b4" }, "outputs": [ { "data": { "text/plain": [ "[{'id': 'doc:5b403537-786d-4ba4-b7e6-5d67383dd3fd:chunk_81',\n", " 'vector_distance': '0.344038009644',\n", " 'doc_id': '[\"5b403537-786d-4ba4-b7e6-5d67383dd3fd\"]',\n", " 'chunk_id': '[\"chunk_81\"]',\n", " 'allowed_roles': '[\"executive\",\"finance\"]',\n", " 'content': '[\"iPhone $ 205,489 $ 191,973 $ 137,781 \\\\nMac 40,177 35,190 28,622 \\\\niPad 29,292 31,862 23,724 \\\\nWearables, Home and Accessories 41,241 38,367 30,620 \\\\nServices 78,129 68,425 53,768 \\\\nTotal net sales $ 394,328 $ 365,817 $ 274,515 \\\\n(1) Products net sales include amortization of the deferred value of unspecified software upgrade rights, which are bundled in the sales price of the respectiveproduct.\\\\n(2) Wearables, Home and Accessories net sales include sales of AirPods, Apple TV, Apple Watch, Beats products, HomePod mini and accessories.\\\\n(3) Services net sales include sales from the Company’s advertising, AppleCare, cloud, digital content, payment and other services. Services net sales also include\\\\namortization of the deferred value of services bundled in the sales price of certain products.\\\\n(4) Includes $7.5 billion of revenue recognized in 2022 that was included in deferred revenue as of September 25, 2021, $6.7 billion of revenue recognized in 2021that was included in deferred revenue as of September 26, 2020, and $5.0 billion of revenue recognized in 2020 that was included in deferred revenue as of\\\\nSeptember 28, 2019.\\\\nThe Company’s proportion of net sales by disaggregated revenue source was generally consistent for each reportable segment in Note 11, “SegmentInformation and Geographic Data” for 2022, 2021 and 2020, except in Greater China, where iPhone revenue represented a moderately higher proportion of net\\\\nsales in 2022 and 2021.\\\\nAs of September 24, 2022 and September 25, 2021, the Company had total deferred revenue of $12.4 billion and $11.9 billion, respectively. As of September 24,\\\\n2022, the Company expects 64% of total deferred revenue to be realized in less than a year, 27% within one-to-two years, 7% within two-to-three years and 2%in greater than three years.\\\\n (1)\\\\n(1)\\\\n (1)\\\\n(1)(2)\\\\n(3)\\\\n(4)\\\\nApple Inc. | 2022 Form 10-K | 37\"]'},\n", " {'id': 'doc:5b403537-786d-4ba4-b7e6-5d67383dd3fd:chunk_68',\n", " 'vector_distance': '0.353463172913',\n", " 'doc_id': '[\"5b403537-786d-4ba4-b7e6-5d67383dd3fd\"]',\n", " 'chunk_id': '[\"chunk_68\"]',\n", " 'allowed_roles': '[\"executive\",\"finance\"]',\n", " 'content': '[\"Apple Inc.\\\\nCONSOLIDATED STATEMENTS OF OPERATIONS(In millions, except number of shares which are reflected in thousands and per share amounts)\\\\nYears ended\\\\nSeptember 24,2022 September 25,2021 September 26,2020\\\\nNet sales:\\\\n Products $ 316,199 $ 297,392 $ 220,747 \\\\n Services 78,129 68,425 53,768 \\\\nTotal net sales 394,328 365,817 274,515 \\\\nCost of sales:\\\\n Products 201,471 192,266 151,286 \\\\n Services 22,075 20,715 18,273 \\\\nTotal cost of sales 223,546 212,981 169,559 \\\\nGross margin 170,782 152,836 104,956 \\\\nOperating expenses:\\\\nResearch and development 26,251 21,914 18,752 \\\\nSelling, general and administrative 25,094 21,973 19,916 \\\\nTotal operating expenses 51,345 43,887 38,668 \\\\nOperating income 119,437 108,949 66,288 \\\\nOther income/(expense), net (334) 258 803 \\\\nIncome before provision for income taxes 119,103 109,207 67,091 \\\\nProvision for income taxes 19,300 14,527 9,680 \\\\nNet income $ 99,803 $ 94,680 $ 57,411 \\\\nEarnings per share:\\\\nBasic $ 6.15 $ 5.67 $ 3.31 \\\\nDiluted $ 6.11 $ 5.61 $ 3.28 \\\\nShares used in computing earnings per share:\\\\nBasic 16,215,963 16,701,272 17,352,119 \\\\nDiluted 16,325,819 16,864,919 17,528,214 \\\\nSee accompanying Notes to Consolidated Financial Statements.\\\\nApple Inc. | 2022 Form 10-K | 29\"]'},\n", " {'id': 'doc:5b403537-786d-4ba4-b7e6-5d67383dd3fd:chunk_72',\n", " 'vector_distance': '0.354508280754',\n", " 'doc_id': '[\"5b403537-786d-4ba4-b7e6-5d67383dd3fd\"]',\n", " 'chunk_id': '[\"chunk_72\"]',\n", " 'allowed_roles': '[\"executive\",\"finance\"]',\n", " 'content': '[\"Apple Inc.\\\\nCONSOLIDATED STATEMENTS OF CASH FLOWS(In millions)\\\\nYears ended\\\\nSeptember 24,2022 September 25,2021 September 26,2020\\\\nCash, cash equivalents and restricted cash, beginning balances $ 35,929 $ 39,789 $ 50,224 \\\\nOperating activities:\\\\nNet income 99,803 94,680 57,411 \\\\nAdjustments to reconcile net income to cash generated by operating activities:\\\\nDepreciation and amortization 11,104 11,284 11,056 \\\\nShare-based compensation expense 9,038 7,906 6,829 \\\\nDeferred income tax expense/(benefit) 895 (4,774) (215)\\\\nOther 111 (147) (97)\\\\nChanges in operating assets and liabilities:\\\\nAccounts receivable, net (1,823) (10,125) 6,917 \\\\nInventories 1,484 (2,642) (127)\\\\nVendor non-trade receivables (7,520) (3,903) 1,553 \\\\nOther current and non-current assets (6,499) (8,042) (9,588)\\\\nAccounts payable 9,448 12,326 (4,062)\\\\nDeferred revenue 478 1,676 2,081 \\\\nOther current and non-current liabilities 5,632 5,799 8,916 \\\\nCash generated by operating activities 122,151 104,038 80,674 \\\\nInvesting activities:\\\\nPurchases of marketable securities (76,923) (109,558) (114,938)\\\\nProceeds from maturities of marketable securities 29,917 59,023 69,918 \\\\nProceeds from sales of marketable securities 37,446 47,460 50,473 \\\\nPayments for acquisition of property, plant and equipment (10,708) (11,085) (7,309)\\\\nPayments made in connection with business acquisitions, net (306) (33) (1,524)\\\\nOther (1,780) (352) (909)\\\\nCash used in investing activities (22,354) (14,545) (4,289)\\\\nFinancing activities:\\\\nPayments for taxes related to net share settlement of equity awards (6,223) (6,556) (3,634)\\\\nPayments for dividends and dividend equivalents (14,841) (14,467) (14,081)\\\\nRepurchases of common stock (89,402) (85,971) (72,358)\\\\nProceeds from issuance of term debt, net 5,465 20,393 16,091 \\\\nRepayments of term debt (9,543) (8,750) (12,629)\\\\nProceeds from/(Repayments of) commercial paper, net 3,955 1,022 (963)\\\\nOther (160) 976 754 \\\\nCash used in financing activities (110,749) (93,353) (86,820)\\\\nDecrease in cash, cash equivalents and restricted cash (10,952) (3,860) (10,435)\\\\nCash, cash equivalents and restricted cash, ending balances $ 24,977 $ 35,929 $ 39,789 \\\\nSupplemental cash flow disclosure:\\\\nCash paid for income taxes, net $ 19,573 $ 25,385 $ 9,501 \\\\nCash paid for interest $ 2,865 $ 2,687 $ 3,002 \\\\nSee accompanying Notes to Consolidated Financial Statements.\\\\nApple Inc. | 2022 Form 10-K | 33\"]'}]" ] }, "execution_count": 24, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Query with valid user\n", "results = user_query(\n", " alice.user_id,\n", " query=\"What was the total revenue amount for Apple according to their 10k?\"\n", ")\n", "results[:3]" ] }, { "cell_type": "markdown", "id": "b3b432e6", "metadata": { "id": "b3b432e6" }, "source": [ "## 5. Implementing Role-Based RAG from scratch\n", "*with OpenAI and Redis*" ] }, { "cell_type": "code", "execution_count": 37, "id": "794b3c41", "metadata": { "id": "794b3c41" }, "outputs": [], "source": [ "from openai import OpenAI\n", "from typing import List, Optional\n", "import os\n", "\n", "from redisvl.extensions.message_history import MessageHistory\n", "\n", "\n", "class RAGChatManager:\n", " \"\"\"\n", " Manages RAG-enhanced chat interactions with role-based access control and chat history.\n", "\n", " Attributes:\n", " kb: A KnowledgeBase instance for searching documents\n", " client: An OpenAI client for chat completions\n", " model: Name of OpenAI model to use\n", " sessions: Dict to store active chat sessions\n", " system_prompt: The default system prompt\n", " \"\"\"\n", "\n", " def __init__(\n", " self,\n", " knowledge_base: \"KnowledgeBase\",\n", " openai_api_key: Optional[str] = None,\n", " openai_model: str = \"gpt-4\",\n", " system_prompt: str = \"You are a helpful chatbot assistant with access to knowledge base documents\"\n", " ):\n", " \"\"\"Initialize the RAG chat manager.\"\"\"\n", " self.kb = knowledge_base\n", " self.client = OpenAI(api_key=openai_api_key or os.getenv(\"OPENAI_API_KEY\"))\n", " self.model = openai_model\n", " self.sessions = {}\n", " self.system_prompt = system_prompt\n", "\n", " def user_roles(self, user_id: str) -> set:\n", " \"\"\"\n", " Get and validate user roles.\n", "\n", " Args:\n", " user_id: User identifier\n", "\n", " Returns:\n", " Set of user roles\n", "\n", " Raises:\n", " ValueError: If user not found or has no roles\n", " \"\"\"\n", " user_obj = User.get(self.kb.redis_client, user_id)\n", " if not user_obj:\n", " raise ValueError(f\"User {user_id} not found.\")\n", "\n", " roles = set([role.value for role in user_obj.roles])\n", " if not roles:\n", " raise ValueError(f\"User {user_id} does not have any roles.\")\n", "\n", " return roles\n", "\n", " def start_session(self, user_id: str) -> None:\n", " \"\"\"\n", " Start a new chat session for a user.\n", "\n", " Args:\n", " user_id: User identifier\n", " \"\"\"\n", " if user_id not in self.sessions:\n", " self.sessions[user_id] = MessageHistory(\n", " name=f\"session:{user_id}\",\n", " redis_client=self.kb.redis_client\n", " )\n", "\n", " def prep_msgs(\n", " self,\n", " user_id: str,\n", " system_prompt: str,\n", " context: str,\n", " query: str\n", " ) -> List[dict]:\n", " \"\"\"\n", " Get chat history messages including system prompt.\n", "\n", " Args:\n", " user_id: User identifier for the session\n", " system_prompt: Optional system prompt to prepend\n", " context: Relevant context fetched from the knowledge base\n", " query: Original user question\n", "\n", " Returns:\n", " List of message dictionaries\n", " \"\"\"\n", " messages = [{\"role\": \"system\", \"content\": system_prompt}]\n", "\n", " if user_id in self.sessions:\n", " messages.extend(self.sessions[user_id].get_recent())\n", "\n", " messages.append({\n", " \"role\": \"user\",\n", " \"content\": f\"\"\"Context information is below.\n", " ---------------------\n", " {context}\n", " ---------------------\n", " Given the context information above and the chat conversation history, please answer the question faithfully: {query}\"\"\"\n", " })\n", "\n", " for msg in messages:\n", " if msg[\"role\"] == \"llm\":\n", " msg[\"role\"] = \"assistant\"\n", "\n", " return messages\n", "\n", " def chat(self, user_id: str, system_prompt: Optional[str] = None) -> None:\n", " \"\"\"\n", " Start an interactive chat loop with the user.\n", "\n", " Args:\n", " user_id: User identifier\n", " system_prompt: Optional system prompt\n", "\n", " The loop continues until user types 'exit' or 'quit'\n", " \"\"\"\n", " self.start_session(user_id)\n", "\n", " print(\"Starting chat session with GPT4. Type 'exit' or 'quit' to end the session.\")\n", " while True:\n", " query = input(\"\\nYou: \").strip()\n", "\n", " if query.lower() in ['exit', 'quit']:\n", " print(\"\\nEnding chat session...\")\n", " break\n", "\n", " response = self.answer(query, user_id, system_prompt)\n", " print(f\"\\nUser: {query}\")\n", " print(f\"Assistant: {response}\")\n", "\n", " def answer(\n", " self,\n", " query: str,\n", " user_id: str,\n", " system_prompt: Optional[str] = None\n", " ) -> str:\n", " \"\"\"\n", " Process a chat message with RAG enhancement and role-based access.\n", "\n", " If any exception occurs at any stage (roles, document search, LLM call),\n", " we do NOT store anything in the session and simply return the error.\n", " Otherwise, we store the query and the response (including 'no docs found' case).\n", "\n", " Args:\n", " query: User's question\n", " user_id: User identifier\n", " system_prompt: Optional system prompt\n", "\n", " Returns:\n", " AI response string or error message\n", " \"\"\"\n", "\n", " # Start or retrieve an existing session for user\n", " self.start_session(user_id)\n", "\n", " try:\n", " # 1. Validate user roles\n", " roles = self.user_roles(user_id)\n", "\n", " # 2. Use provided system prompt or default\n", " system_prompt = system_prompt or self.system_prompt\n", "\n", " # 3. Search for relevant documents\n", " docs = self.kb.search(query, roles)\n", "\n", " # 4. If no documents, store & return early\n", " if not docs:\n", " no_docs_msg = (\n", " \"I couldn't find any relevant documents you have permission to access. \"\n", " \"Please try rephrasing your question or contact an administrator if you believe this is an error.\"\n", " )\n", " self.sessions[user_id].store(query, no_docs_msg)\n", " return no_docs_msg\n", "\n", " # 5. Prepare context and messages for the LLM\n", " context = \"\\n\\n\".join([doc.get(\"content\", \"\") for doc in docs])\n", " messages = self.prep_msgs(\n", " user_id=user_id,\n", " system_prompt=system_prompt,\n", " context=context,\n", " query=query\n", " )\n", "\n", " # 6. Generate response from the model\n", " response = self.client.chat.completions.create(\n", " model=self.model,\n", " messages=messages\n", " )\n", " ai_response = response.choices[0].message.content\n", "\n", " # 7. Store query and LLM response\n", " self.sessions[user_id].store(query, ai_response)\n", "\n", " return ai_response\n", "\n", " except Exception as e:\n", " # Catch any exception; do not store anything, just return the error.\n", " return f\"I encountered an error: {str(e)}\"\n" ] }, { "cell_type": "markdown", "id": "zJdHMGdUCl_S", "metadata": { "id": "zJdHMGdUCl_S" }, "source": [ "### Session-aware, role-based RAG" ] }, { "cell_type": "code", "execution_count": 32, "id": "1HDy2Ltr12I1", "metadata": { "id": "1HDy2Ltr12I1" }, "outputs": [], "source": [ "bot = RAGChatManager(kb)" ] }, { "cell_type": "code", "execution_count": 33, "id": "sM6BQ-ZL2LUf", "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 89 }, "id": "sM6BQ-ZL2LUf", "outputId": "b678b1ac-e177-4d16-9af8-2cd2cf2e48c1" }, "outputs": [ { "data": { "text/plain": [ "\"The context information and the chat conversation history provided does not contain any details on a vehicle's make and model.\"" ] }, "execution_count": 33, "metadata": {}, "output_type": "execute_result" } ], "source": [ "bot.answer(\"What is the make and model of the vehicle?\", user_id=\"alice\")" ] }, { "cell_type": "code", "execution_count": 28, "id": "3iJdgsaAjsaA", "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 89 }, "id": "3iJdgsaAjsaA", "outputId": "545b9621-e04e-4d96-ade7-5ad1e1311d3c" }, "outputs": [ { "data": { "text/plain": [ "'The make and model of the vehicle is the Chevrolet Colorado.'" ] }, "execution_count": 28, "metadata": {}, "output_type": "execute_result" } ], "source": [ "bot.answer(\"What is the make and model of the vehicle?\", user_id=\"tyler\")" ] }, { "cell_type": "code", "execution_count": 29, "id": "17CUi5TXBFSB", "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 71 }, "id": "17CUi5TXBFSB", "outputId": "852635cc-01a4-4a02-d07d-4a48eabafbba" }, "outputs": [ { "data": { "text/plain": [ "'The vehicle is from the year 2022.'" ] }, "execution_count": 29, "metadata": {}, "output_type": "execute_result" } ], "source": [ "bot.answer(\"What year is it?\", user_id=\"tyler\")" ] }, { "cell_type": "code", "execution_count": 38, "id": "N4IV1bLTCj1N", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "N4IV1bLTCj1N", "outputId": "e456deb7-c15d-4a88-ad31-27782be58f72" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Starting chat session with GPT4. Type 'exit' or 'quit' to end the session.\n", "\n", "User: What is the towing capacity of the truck?\n", "\n", "Assistant: The towing capacity of the truck depends on the specific model and configurations:\n", "\n", "- For models with the 2.5L DOHC I-4 engine, the maximum trailering weight rating is 3,500 lbs.\n", "- For models with the 3.6L DOHC V6 engine, the maximum trailering weight rating is 7,000 lbs.\n", "- For models with the Duramax 2.8L Turbo-Diesel I-4 engine, the maximum trailering weight rating is 7,700 lbs.\n", "- For the ZR2 models, regardless of engine, the towing capacity is up to 5,000 lbs.\n", "\n", "These figures are intended for comparison purposes only. Before using the vehicle for trailering, you should carefully review the Trailering section of the Owner’s Manual. The trailering capacity of your specific vehicle may vary. The weight of passengers, cargo, and options or accessories may reduce the amount you can tow.\n", "\n", "User: Is it generally safe to drive? What safety features are available?\n", "\n", "Assistant: Yes, it is generally safe to drive the vehicle. It comes equipped with several safety features including:\n", "\n", "1. Electronic Stability Control System and Traction Control: This includes Hill Start Assist that helps maintain stability by adjusting the brakes and engine power during certain driving conditions.\n", "\n", "2. Airbags: The vehicle comes with dual-stage frontal airbags for the driver and front passenger, seat-mounted side-impact airbags for the driver and front passenger, and head-curtain airbags for front and rear outboard seating positions.\n", "\n", "3. Teen Driver technology: This feature allows parents to monitor and manage their teenagers' driving behavior by limiting certain features and promoting safer driving habits.\n", "\n", "4. Tire Pressure Monitoring System: This feature helps maintain proper tire pressure, improving safety and fuel efficiency.\n", "\n", "5. Rear Vision Camera: This feature provides a clear view of the area behind the vehicle, improving safety when reversing.\n", "\n", "6. Hitch Guidance with dynamic trailering assist guideline: This feature significantly improves safety when the truck is hauling loads.\n", "\n", "Safety on the road also depends largely on the driver — it's crucial to operate the vehicle safely, and to be aware and respectful of road and traffic conditions.\n", "\n", "Ending chat session...\n" ] } ], "source": [ "# NBVAL_SKIP\n", "bot.chat(user_id=\"tyler\")" ] }, { "cell_type": "markdown", "id": "SHg3tFa2u0Nh", "metadata": { "id": "SHg3tFa2u0Nh" }, "source": [ "## 6. Summary & Next Steps\n", "\n", "In this notebook, we set up a **basic** for a Role-Based RAG system:\n", "\n", "1. **Users** (with `roles`) stored in Redis via JSON.\n", "2. **Documents** (with `allowed_roles`) loaded, parsed, embedded and also stored in Redis.\n", "3. A user search pipeline that honors user roles when retrieving documents.\n", "\n", "\n", "This approach ensures that **only documents** whose roles match the user’s roles are returned.\n", "\n", "\n", "With these building blocks in place, you can integrate an LLM to supply a context from the returned docs, producing a robust retrieval-augmented generation pipeline with role-based access controls.\n" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4", "provenance": [] }, "kernelspec": { "display_name": ".venv (3.12.5)", "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.5" } }, "nbformat": 4, "nbformat_minor": 5 }