Multi-Tenant RAG: Namespace Isolation & Security Guide

-- Enable RLS on document chunks table
CREATE TABLE document_chunks (
    id BIGSERIAL PRIMARY KEY,
    tenant_id TEXT NOT NULL,
    content TEXT NOT NULL,
    metadata JSONB,
    embedding vector(1536),
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- Enable Row-Level Security
ALTER TABLE document_chunks ENABLE ROW LEVEL SECURITY;

-- Policy: Users can only see chunks from their tenant
CREATE POLICY tenant_isolation_policy ON document_chunks
    FOR ALL
    USING (tenant_id = current_setting('app.current_tenant_id', true));

-- Create index on tenant_id for filtering performance
CREATE INDEX idx_chunks_tenant ON document_chunks(tenant_id);
CREATE INDEX idx_chunks_tenant_embedding ON document_chunks USING ivfflat (embedding vector_cosine_ops)
    WHERE tenant_id = current_setting('app.current_tenant_id', true);

import asyncpg
from pgvector.asyncpg import register_vector
from openai import OpenAI

client = OpenAI()

class MultiTenantRAG:
    def __init__(self, db_pool: asyncpg.Pool):
        self.db_pool = db_pool

    async def query(self, question: str, tenant_id: str, top_k: int = 5) -> dict:
        """Execute multi-tenant RAG query with RLS enforcement."""
        async with self.db_pool.acquire() as conn:
            await conn.execute("SET app.current_tenant_id = $1", tenant_id)

            try:
                # Embed query
                query_embedding = client.embeddings.create(
                    input=[question],
                    model="text-embedding-3-small"
                ).data[0].embedding

                # Retrieve chunks (RLS automatically filters by tenant_id)
                await register_vector(conn)
                chunks = await conn.fetch(
                    """
                    SELECT content, metadata, 1 - (embedding <=> $1) AS similarity
                    FROM document_chunks
                    ORDER BY embedding <=> $1
                    LIMIT $2
                    """,
                    query_embedding,
                    top_k
                )

                # Generate answer
                context = "\n\n---\n\n".join(c["content"] for c in chunks)
                response = client.chat.completions.create(
                    model="gpt-4o",
                    messages=[
                        {"role": "system", "content": f"Answer for tenant {tenant_id}. Use only provided context."},
                        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
                    ],
                    temperature=0.1
                )

                return {
                    "answer": response.choices[0].message.content,
                    "sources": [c["metadata"].get("source") for c in chunks],
                    "tenant_id": tenant_id,
                }

            finally:
                # Reset tenant context
                await conn.execute("RESET app.current_tenant_id")

    async def ingest_document(self, doc: dict, tenant_id: str):
        """Ingest document for specific tenant with RLS."""
        async with self.db_pool.acquire() as conn:
            await conn.execute("SET app.current_tenant_id = $1", tenant_id)

            try:
                # Chunk and embed
                chunks = self._chunk_document(doc)
                embeddings = self._embed_chunks(chunks)

                # Insert with tenant_id (RLS will validate)
                for chunk, embedding in zip(chunks, embeddings):
                    await register_vector(conn)
                    await conn.execute(
                        """
                        INSERT INTO document_chunks (tenant_id, content, metadata, embedding)
                        VALUES ($1, $2, $3, $4)
                        """,
                        tenant_id,
                        chunk["content"],
                        chunk["metadata"],
                        embedding
                    )

            finally:
                await conn.execute("RESET app.current_tenant_id")

from pinecone import Pinecone

pc = Pinecone(api_key="...")
index = pc.Index("multi-tenant-rag")

class PineconeMultiTenantRAG:
    def __init__(self, index):
        self.index = index

    async def upsert_chunks(self, chunks: list[dict], embeddings: list[list[float]], tenant_id: str):
        """Upsert chunks to tenant-specific namespace."""
        vectors = [
            {
                "id": f"{tenant_id}_{i}",
                "values": embedding,
                "metadata": {**chunk["metadata"], "tenant_id": tenant_id}
            }
            for i, (chunk, embedding) in enumerate(zip(chunks, embeddings))
        ]

        # Use namespace for isolation
        self.index.upsert(vectors=vectors, namespace=tenant_id)

    async def query(self, query_embedding: list[float], tenant_id: str, top_k: int = 5) -> list[dict]:
        """Query tenant-specific namespace."""
        results = self.index.query(
            vector=query_embedding,
            top_k=top_k,
            namespace=tenant_id,  # Isolates to tenant namespace
            include_metadata=True
        )

        return [
            {
                "content": match["metadata"].get("content"),
                "score": match["score"],
                "metadata": match["metadata"]
            }
            for match in results["matches"]
        ]

    async def delete_tenant_data(self, tenant_id: str):
        """Hard delete all tenant data."""
        # Delete entire namespace
        self.index.delete(namespace=tenant_id, delete_all=True)

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

client = QdrantClient(url="http://localhost:6333")

class QdrantMultiTenantRAG:
    def __init__(self, client: QdrantClient):
        self.client = client

    async def create_tenant(self, tenant_id: str):
        """Create dedicated collection for tenant."""
        collection_name = f"tenant_{tenant_id}"

        self.client.create_collection(
            collection_name=collection_name,
            vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
        )

    async def upsert_chunks(self, chunks: list[dict], embeddings: list[list[float]], tenant_id: str):
        """Upsert to tenant-specific collection."""
        collection_name = f"tenant_{tenant_id}"

        points = [
            PointStruct(
                id=i,
                vector=embedding,
                payload={"content": chunk["content"], **chunk["metadata"]}
            )
            for i, (chunk, embedding) in enumerate(zip(chunks, embeddings))
        ]

        self.client.upsert(collection_name=collection_name, points=points)

    async def query(self, query_embedding: list[float], tenant_id: str, top_k: int = 5) -> list[dict]:
        """Query tenant-specific collection."""
        collection_name = f"tenant_{tenant_id}"

        results = self.client.search(
            collection_name=collection_name,
            query_vector=query_embedding,
            limit=top_k
        )

        return [
            {
                "content": hit.payload.get("content"),
                "score": hit.score,
                "metadata": hit.payload
            }
            for hit in results
        ]

    async def delete_tenant_data(self, tenant_id: str):
        """Hard delete tenant collection."""
        collection_name = f"tenant_{tenant_id}"
        self.client.delete_collection(collection_name=collection_name)

# Weaviate metadata filtering example
import weaviate

client = weaviate.Client(url="http://localhost:8080")

async def query_with_filter(query_embedding: list[float], tenant_id: str, top_k: int = 5):
    """Query with tenant_id metadata filter."""
    results = client.query.get("Document", ["content", "metadata"]).with_near_vector({
        "vector": query_embedding
    }).with_where({
        "path": ["tenant_id"],
        "operator": "Equal",
        "valueText": tenant_id
    }).with_limit(top_k).do()

    return results["data"]["Get"]["Document"]

import hashlib
from typing import Optional

class SecureEmbeddingPipeline:
    def __init__(self, embedding_client):
        self.client = embedding_client
        self.tenant_keys = {}  # Store tenant encryption keys

    async def embed_with_tenant_context(self, texts: list[str], tenant_id: str) -> list[list[float]]:
        """Embed texts with tenant-specific context marker."""
        # Option 1: Prepend tenant marker (subtle bias)
        marked_texts = [f"[Tenant: {tenant_id}] {text}" for text in texts]

        embeddings = self.client.embeddings.create(
            input=marked_texts,
            model="text-embedding-3-small"
        )

        return [item.embedding for item in embeddings.data]

    async def embed_with_audit(self, texts: list[str], tenant_id: str, user_id: Optional[str] = None) -> tuple:
        """Embed with audit logging."""
        embeddings = await self.embed_with_tenant_context(texts, tenant_id)

        # Log embedding operation
        await self._audit_log({
            "operation": "embed",
            "tenant_id": tenant_id,
            "user_id": user_id,
            "text_count": len(texts),
            "timestamp": "...",
        })

        return embeddings

    def _hash_content(self, content: str) -> str:
        """Hash content for deduplication and privacy."""
        return hashlib.sha256(content.encode()).hexdigest()

class AuditedMultiTenantRAG:
    def __init__(self, rag_system, audit_logger):
        self.rag = rag_system
        self.audit = audit_logger

    async def query_with_audit(self, question: str, tenant_id: str, user_id: str) -> dict:
        """Query with comprehensive audit logging."""
        query_id = str(uuid.uuid4())

        # Log query attempt
        await self.audit.log({
            "query_id": query_id,
            "event": "query_start",
            "tenant_id": tenant_id,
            "user_id": user_id,
            "question_hash": hashlib.sha256(question.encode()).hexdigest(),
            "timestamp": datetime.utcnow(),
        })

        try:
            # Execute query
            result = await self.rag.query(question, tenant_id)

            # Log successful retrieval
            await self.audit.log({
                "query_id": query_id,
                "event": "retrieval_success",
                "tenant_id": tenant_id,
                "chunks_retrieved": len(result.get("sources", [])),
                "sources": result.get("sources", []),
                "timestamp": datetime.utcnow(),
            })

            return result

        except Exception as e:
            # Log failure
            await self.audit.log({
                "query_id": query_id,
                "event": "query_error",
                "tenant_id": tenant_id,
                "error": str(e),
                "timestamp": datetime.utcnow(),
            })
            raise

import asyncio
from typing import Optional
from dataclasses import dataclass

@dataclass
class TenantConfig:
    tenant_id: str
    embedding_model: str = "text-embedding-3-small"
    chunk_size: int = 512
    max_chunks_per_query: int = 5
    enable_audit: bool = True

class ProductionMultiTenantRAG:
    def __init__(self, db_pool, vector_store, audit_logger):
        self.db_pool = db_pool
        self.vector_store = vector_store
        self.audit = audit_logger
        self.tenant_configs = {}

    async def register_tenant(self, config: TenantConfig):
        """Register new tenant with dedicated resources."""
        # Create database schema (if using schema-per-tenant)
        async with self.db_pool.acquire() as conn:
            await conn.execute("SET app.current_tenant_id = $1", config.tenant_id)
            # Create tenant-specific tables if needed

        # Create vector store namespace
        await self.vector_store.create_namespace(config.tenant_id)

        # Store config
        self.tenant_configs[config.tenant_id] = config

        await self.audit.log({
            "event": "tenant_registered",
            "tenant_id": config.tenant_id,
            "timestamp": datetime.utcnow(),
        })

    async def ingest(self, doc: dict, tenant_id: str, user_id: Optional[str] = None):
        """Ingest document with full isolation."""
        config = self.tenant_configs.get(tenant_id)
        if not config:
            raise ValueError(f"Tenant {tenant_id} not registered")

        # Chunk
        chunks = self._chunk_document(doc, config.chunk_size)

        # Embed with tenant context
        embeddings = await self._embed_with_context(
            [c["content"] for c in chunks],
            tenant_id,
            config.embedding_model
        )

        # Store in database with RLS
        async with self.db_pool.acquire() as conn:
            await conn.execute("SET app.current_tenant_id = $1", tenant_id)
            try:
                for chunk, embedding in zip(chunks, embeddings):
                    await conn.execute(
                        """
                        INSERT INTO document_chunks (tenant_id, content, metadata, embedding)
                        VALUES ($1, $2, $3, $4)
                        """,
                        tenant_id, chunk["content"], chunk["metadata"], embedding
                    )
            finally:
                await conn.execute("RESET app.current_tenant_id")

        # Store in vector database with namespace isolation
        await self.vector_store.upsert(chunks, embeddings, namespace=tenant_id)

        if config.enable_audit:
            await self.audit.log({
                "event": "document_ingested",
                "tenant_id": tenant_id,
                "user_id": user_id,
                "doc_id": doc.get("id"),
                "chunks_created": len(chunks),
            })

    async def query(self, question: str, tenant_id: str, user_id: str) -> dict:
        """Query with full isolation and audit."""
        config = self.tenant_configs.get(tenant_id)
        if not config:
            raise ValueError(f"Tenant {tenant_id} not registered")

        query_id = str(uuid.uuid4())

        if config.enable_audit:
            await self.audit.log({
                "query_id": query_id,
                "event": "query_start",
                "tenant_id": tenant_id,
                "user_id": user_id,
            })

        # Embed query
        query_embedding = await self._embed_with_context([question], tenant_id, config.embedding_model)

        # Retrieve from tenant namespace only
        chunks = await self.vector_store.query(
            query_embedding[0],
            namespace=tenant_id,
            top_k=config.max_chunks_per_query
        )

        # Generate answer
        context = "\n\n---\n\n".join(c["content"] for c in chunks)
        answer = await self._generate_answer(question, context, tenant_id)

        if config.enable_audit:
            await self.audit.log({
                "query_id": query_id,
                "event": "query_complete",
                "tenant_id": tenant_id,
                "chunks_retrieved": len(chunks),
            })

        return {"answer": answer, "sources": chunks, "query_id": query_id}

    async def delete_tenant(self, tenant_id: str):
        """Hard delete all tenant data."""
        # Delete from database
        async with self.db_pool.acquire() as conn:
            await conn.execute("DELETE FROM document_chunks WHERE tenant_id = $1", tenant_id)

        # Delete from vector store
        await self.vector_store.delete_namespace(tenant_id)

        # Remove config
        self.tenant_configs.pop(tenant_id, None)

        await self.audit.log({
            "event": "tenant_deleted",
            "tenant_id": tenant_id,
            "timestamp": datetime.utcnow(),
        })