RAG Pipeline Explained — The Embedding Drift That Cost Us $4k/Month in Token Waste

import chromadb
from sentence_transformers import SentenceTransformer
from typing import List, Dict
import logging

logger = logging.getLogger(__name__)

class RAGRetriever:
    def __init__(self, collection_name: str = "docs", persist_dir: str = "./chroma"):
        # Pin the exact model revision to prevent silent drift
        self.model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
        self.client = chromadb.PersistentClient(path=persist_dir)
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"}  # Use cosine distance, not L2
        )
        logger.info(f"Connected to collection '{collection_name}' with {self.collection.count()} documents")

    def embed_query(self, query: str) -> List[float]:
        """Embed a query string. Returns a list of floats."""
        embedding = self.model.encode(query, normalize_embeddings=True).tolist()
        return embedding

    def retrieve(self, query: str, top_k: int = 5) -> List[Dict]:
        """Retrieve top-k documents for a query."""
        query_embedding = self.embed_query(query)
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            include=["documents", "metadatas", "distances"]
        )
        # Re-rank: put the most relevant chunk last (recency effect)
        # Chroma returns results sorted by distance (ascending), so index 0 is closest
        docs = []
        for i in range(len(results['ids'][0])):
            docs.append({
                'id': results['ids'][0][i],
                'text': results['documents'][0][i],
                'metadata': results['metadatas'][0][i],
                'score': 1 - results['distances'][0][i]  # Convert distance to similarity
            })
        # Sort by score ascending so highest score is last in the prompt
        docs.sort(key=lambda x: x['score'])
        return docs

    def format_context(self, docs: List[Dict]) -> str:
        """Format retrieved docs into a context string."""
        # Put the highest-relevance chunk last, near the query
        context_parts = []
        for doc in docs:
            context_parts.append(f"[Source: {doc['metadata'].get('source', 'unknown')}]\n{doc['text']}")
        return "\n\n---\n\n".join(context_parts)

import hashlib
from typing import List, Optional
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, WebBaseLoader
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
import chromadb
import logging

logger = logging.getLogger(__name__)

class ProductionRAGPipeline:
    def __init__(self, persist_directory: str = "./chroma_db"):
        # Use OpenAI's cheapest embedding model
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
        self.persist_directory = persist_directory
        # Use a persistent client to avoid reloading the index every time
        self.vector_store = Chroma(
            collection_name="rag_docs",
            embedding_function=self.embeddings,
            persist_directory=self.persist_directory
        )
        # Chunk size of 512 tokens with 64-token overlap
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=512,
            chunk_overlap=64,
            separators=["\n\n", "\n", ".", " ", ""],
            length_function=len  # Approximate; use tiktoken for exact token count
        )
        logger.info(f"Initialized RAG pipeline with {self.vector_store._collection.count()} existing docs")

    def _content_hash(self, text: str) -> str:
        """Compute a hash of the text to detect changes."""
        return hashlib.sha256(text.encode('utf-8')).hexdigest()

    def load_pdf(self, file_path: str) -> List[dict]:
        """Load and chunk a PDF file."""
        loader = PyPDFLoader(file_path)
        documents = loader.load()
        chunks = self.text_splitter.split_documents(documents)
        # Add metadata: source and content hash
        for chunk in chunks:
            chunk.metadata['source'] = file_path
            chunk.metadata['content_hash'] = self._content_hash(chunk.page_content)
        return chunks

    def load_webpage(self, url: str) -> List[dict]:
        """Load and chunk a webpage."""
        loader = WebBaseLoader(url)
        documents = loader.load()
        chunks = self.text_splitter.split_documents(documents)
        for chunk in chunks:
            chunk.metadata['source'] = url
            chunk.metadata['content_hash'] = self._content_hash(chunk.page_content)
        return chunks

    def index_documents(self, chunks: List[dict]) -> int:
        """Index chunks into vector store, skipping unchanged content."""
        existing_hashes = set()
        # Fetch all existing hashes from the store (expensive, but necessary for dedup)
        all_metadatas = self.vector_store._collection.get(include=["metadatas"])['metadatas']
        for meta in all_metadatas:
            if 'content_hash' in meta:
                existing_hashes.add(meta['content_hash'])

        new_chunks = []
        for chunk in chunks:
            if chunk.metadata['content_hash'] not in existing_hashes:
                new_chunks.append(chunk)

        if new_chunks:
            self.vector_store.add_documents(new_chunks)
            logger.info(f"Indexed {len(new_chunks)} new chunks")
        else:
            logger.info("No new chunks to index")
        return len(new_chunks)

    def retrieve(self, query: str, k: int = 5, filter: Optional[dict] = None) -> List[dict]:
        """Retrieve top-k chunks for a query, with optional metadata filter."""
        results = self.vector_store.similarity_search_with_relevance_scores(
            query, k=k, filter=filter
        )
        # Re-rank by score ascending (highest last)
        results.sort(key=lambda x: x[1])
        return [{'text': doc.page_content, 'metadata': doc.metadata, 'score': score}
                for doc, score in results]

# Example: when a simple SQL query is better than RAG
import sqlite3
from typing import Optional

class AccountLookup:
    """Use this instead of RAG for structured queries."""
    def __init__(self, db_path: str = "accounts.db"):
        self.conn = sqlite3.connect(db_path)
        self.conn.row_factory = sqlite3.Row

    def get_balance(self, account_id: str) -> Optional[float]:
        cursor = self.conn.execute("SELECT balance FROM accounts WHERE id = ?", (account_id,))
        row = cursor.fetchone()
        return row['balance'] if row else None

# Usage: deterministic, 5ms latency, zero token cost
lookup = AccountLookup()
print(lookup.get_balance("ACC-12345"))  # Returns exact value or None

import redis
import json
import hashlib
from typing import List, Optional

class CachedRAGRetriever:
    def __init__(self, vector_store, redis_host: str = "localhost", redis_port: int = 6379):
        self.vector_store = vector_store
        self.redis = redis.Redis(host=redis_host, port=redis_port, decode_responses=True)
        self.embedding_cache_ttl = 3600  # 1 hour
        self.result_cache_ttl = 300      # 5 minutes

    def _query_hash(self, query: str, filter: Optional[dict] = None) -> str:
        """Create a hash key for caching."""
        key = f"{query}:{json.dumps(filter, sort_keys=True) if filter else 'none'}"
        return hashlib.sha256(key.encode()).hexdigest()

    def retrieve(self, query: str, k: int = 5, filter: Optional[dict] = None) -> List[dict]:
        cache_key = self._query_hash(query, filter)
        # Check result cache first
        cached = self.redis.get(f"result:{cache_key}")
        if cached:
            return json.loads(cached)

        # Check embedding cache
        embed_key = f"embed:{cache_key}"
        cached_embed = self.redis.get(embed_key)
        if cached_embed:
            query_embedding = json.loads(cached_embed)
        else:
            query_embedding = self.vector_store.embeddings.embed_query(query)
            self.redis.setex(embed_key, self.embedding_cache_ttl, json.dumps(query_embedding))

        # Retrieve from vector store
        results = self.vector_store.similarity_search_with_relevance_scores(
            query, k=k, filter=filter
        )
        # Format and cache
        formatted = [{'text': doc.page_content, 'score': score} for doc, score in results]
        self.redis.setex(f"result:{cache_key}", self.result_cache_ttl, json.dumps(formatted))
        return formatted

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
import chromadb

# Mistake 1 fix: Use 512-token chunks with overlap
splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=64,
    separators=["\n\n", "\n", ".", " ", ""]
)

# Mistake 2 fix: Always include metadata filter in retrieval
# When indexing:
# chunk.metadata['product_id'] = 'product_b'
# When querying:
results = vector_store.similarity_search(
    "refund policy",
    k=5,
    filter={"product_id": "product_b"}
)

# Mistake 3 fix: Validate embedding dimensions before indexing
def validate_embedding_dimension(embeddings_model, expected_dim: int = 512):
    """Check that the embedding model outputs the expected dimension."""
    test_embedding = embeddings_model.embed_query("test")
    actual_dim = len(test_embedding)
    if actual_dim != expected_dim:
        raise ValueError(f"Expected {expected_dim} dimensions, got {actual_dim}. Re-index required.")
    print(f"Embedding dimension validated: {actual_dim}")

# Usage
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
validate_embedding_dimension(embeddings, expected_dim=512)

# RAG approach: retrieve and generate
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
prompt = ChatPromptTemplate.from_messages([
    ("system", "Answer the question based on the context provided. Cite the source."),
    ("human", "Context: {context}\n\nQuestion: {question}")
])

# Fine-tuning approach: use a fine-tuned model directly
# (Assuming you have a fine-tuned model ID)
finetuned_llm = ChatOpenAI(model="ft:gpt-4o-mini:your-org:your-model:hash", temperature=0)
# No context needed — the knowledge is in the weights
response = finetuned_llm.invoke("What is the refund policy for Product B?")

import logging
import time
from prometheus_client import Histogram, Counter, Gauge

# Prometheus metrics
RETRIEVAL_LATENCY = Histogram('rag_retrieval_latency_seconds', 'Retrieval latency', buckets=[0.1, 0.25, 0.5, 1.0, 2.0, 5.0])
GENERATION_LATENCY = Histogram('rag_generation_latency_seconds', 'Generation latency', buckets=[0.5, 1.0, 2.0, 5.0, 10.0])
EMPTY_RETRIEVAL = Counter('rag_empty_retrieval_total', 'Number of queries with 0 retrieved documents')
AVG_SIMILARITY = Gauge('rag_avg_similarity', 'Average cosine similarity of retrieved docs')

# Structured logging setup
logger = logging.getLogger(__name__)
handler = logging.StreamHandler()
formatter = logging.Formatter('{"time": "%(asctime)s", "level": "%(levelname)s", "message": %(message)s}')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)

def monitored_retrieve(query: str, retriever):
    """Retrieve with monitoring."""
    start = time.time()
    docs = retriever.retrieve(query)
    duration = time.time() - start
    RETRIEVAL_LATENCY.observe(duration)

    if len(docs) == 0:
        EMPTY_RETRIEVAL.inc()
        logger.warning(json.dumps({"event": "empty_retrieval", "query": query}))

    if docs:
        avg_score = sum(d['score'] for d in docs) / len(docs)
        AVG_SIMILARITY.set(avg_score)
        logger.info(json.dumps({
            "event": "retrieval",
            "query": query,
            "num_docs": len(docs),
            "avg_score": round(avg_score, 3),
            "latency": round(duration, 3)
        }))

    return docs