Advanced RAG: Implementing Hybrid Search with pgvector and RRF

graph TD
    A[User Query] --> B{Query Pre-processing};
    B --> C{Parallel Search Execution};
    C --> D[Dense Vector Search (pgvector)];
    C --> E[Sparse Keyword Search (Postgres FTS)];
    D --> F{Reciprocal Rank Fusion (RRF)};
    E --> F;
    F --> G[Context Re-ranking & Assembly];
    G --> H[Augmented Prompt Generation];
    H --> I[LLM Inference];
    I --> J[Final Response];

CREATE EXTENSION IF NOT EXISTS vector;

-- Our core table for storing document chunks and their embeddings
CREATE TABLE document_chunks (
    id SERIAL PRIMARY KEY,
    document_id INT NOT NULL, -- Foreign key to a parent documents table
    content TEXT NOT NULL,
    embedding VECTOR(384) NOT NULL, -- Using all-MiniLM-L6-v2 which has 384 dimensions
    content_tsv TSVECTOR -- For Full-Text Search
);

-- Create an index on the tsvector column for fast keyword search
CREATE INDEX idx_content_tsv ON document_chunks USING GIN(content_tsv);

-- Create an index for efficient vector similarity search.
-- We'll discuss HNSW vs. IVFFlat later.
-- For now, let's plan for an HNSW index.
CREATE INDEX idx_embedding_hnsw ON document_chunks USING HNSW (embedding vector_cosine_ops);
-- NOTE: HNSW index creation can be slow and memory-intensive. It's often done after the initial data load.

# requirements.txt
psycopg2-binary
sentence-transformers
langchain # For its text splitters
python-dotenv

-- This function will be called by the trigger to update the tsvector column
CREATE OR REPLACE FUNCTION update_content_tsv()
RETURNS TRIGGER AS $$
BEGIN
    NEW.content_tsv := to_tsvector('english', NEW.content);
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- The trigger itself, which fires before any insert or update
CREATE TRIGGER trg_update_content_tsv
BEFORE INSERT OR UPDATE ON document_chunks
FOR EACH ROW
EXECUTE FUNCTION update_content_tsv();

import os
import psycopg2
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter

load_dotenv()

# --- Configuration ---
DB_CONNECT_STR = os.getenv("DB_CONNECT_STR")
MODEL_NAME = 'all-MiniLM-L6-v2' # 384 dimensions

# --- Sample Document ---
# In a real application, this would come from a file, API, etc.
sample_document_text = """
Project Zenith: Q3 2024 Performance Review

The primary performance metric, user engagement, saw a 15% quarter-over-quarter increase. This was largely driven by the rollout of feature 'Nova'. 
However, we encountered a critical production issue tracked as INC-98765. The root cause was identified as a memory leak in the authentication service, specifically related to the JWT caching layer. The error code associated with this failure was ERR_AUTH_CACHE_OOM. A hotfix was deployed on October 3rd.

Financials: Revenue is up 10%, but operational costs increased by 20% due to unscheduled cloud resource scaling to mitigate the INC-98765 incident. The specific product SKU affected was ZNT-PRO-TIER-V2.

Next quarter's focus will be on platform stability and reducing technical debt.
"""

def get_db_connection():
    return psycopg2.connect(DB_CONNECT_STR)

def initialize_model():
    return SentenceTransformer(MODEL_NAME)

def ingest_document(document_text: str, document_id: int, model, conn):
    # 1. Advanced Chunking Strategy
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500, 
        chunk_overlap=50,
        length_function=len,
    )
    chunks = text_splitter.split_text(document_text)
    print(f"Document split into {len(chunks)} chunks.")

    # 2. Generate Embeddings
    embeddings = model.encode(chunks, show_progress_bar=True)

    # 3. Store in PostgreSQL
    with conn.cursor() as cur:
        for i, chunk in enumerate(chunks):
            embedding = embeddings[i].tolist()
            print(f"Inserting chunk {i+1}/{len(chunks)}...")
            cur.execute(
                "INSERT INTO document_chunks (document_id, content, embedding) VALUES (%s, %s, %s)",
                (document_id, chunk, embedding)
            )
    conn.commit()
    print("Ingestion complete.")

if __name__ == "__main__":
    model = initialize_model()
    try:
        conn = get_db_connection()
        ingest_document(sample_document_text, 1, model, conn)
    except Exception as e:
        print(f"An error occurred: {e}")
    finally:
        if 'conn' in locals() and conn:
            conn.close()

def search_dense(query_embedding, top_k: int, conn) -> list[tuple[int, float]]:
    """Performs dense vector search and returns (id, score) tuples."""
    embedding_list = query_embedding.tolist()
    with conn.cursor() as cur:
        cur.execute(
            """SELECT id, 1 - (embedding <=> %s) AS similarity_score
               FROM document_chunks
               ORDER BY embedding <=> %s
               LIMIT %s""",
            (embedding_list, embedding_list, top_k)
        )
        results = cur.fetchall()
    return [(row[0], row[1]) for row in results]

def search_sparse(query_text: str, top_k: int, conn) -> list[tuple[int, float]]:
    """Performs sparse keyword search and returns (id, score) tuples."""
    with conn.cursor() as cur:
        cur.execute(
            """SELECT id, ts_rank(content_tsv, plainto_tsquery('english', %s)) AS rank_score
               FROM document_chunks
               WHERE content_tsv @@ plainto_tsquery('english', %s)
               ORDER BY rank_score DESC
               LIMIT %s""",
            (query_text, query_text, top_k)
        )
        results = cur.fetchall()
    return [(row[0], row[1]) for row in results]

def reciprocal_rank_fusion(list_of_results: list[list[tuple[int, float]]], k: int = 60) -> dict[int, float]:
    """Performs RRF on multiple ranked lists of (id, score) tuples."""
    rrf_scores = {}

    # Each 'results' is a list of (doc_id, score) from one search method
    for results in list_of_results:
        # Enumerate gives us the 0-based index, so we add 1 for 1-based rank
        for rank, (doc_id, _) in enumerate(results, 1):
            if doc_id not in rrf_scores:
                rrf_scores[doc_id] = 0.0
            rrf_scores[doc_id] += 1.0 / (k + rank)

    return rrf_scores

import concurrent.futures

# ... (previous function definitions: search_dense, search_sparse, rrf, etc.)

def hybrid_search(query_text: str, top_k: int, model, conn) -> list[tuple[int, float]]:
    # 1. Generate query embedding
    query_embedding = model.encode(query_text)

    # 2. Execute searches in parallel
    with concurrent.futures.ThreadPoolExecutor() as executor:
        future_dense = executor.submit(search_dense, query_embedding, top_k, conn)
        future_sparse = executor.submit(search_sparse, query_text, top_k, conn)

        dense_results = future_dense.result()
        sparse_results = future_sparse.result()

    print("Dense Results:", dense_results)
    print("Sparse Results:", sparse_results)

    # 3. Fuse results with RRF
    rrf_scores = reciprocal_rank_fusion([dense_results, sparse_results])

    # 4. Sort documents by their final RRF score
    sorted_docs = sorted(rrf_scores.items(), key=lambda item: item[1], reverse=True)

    return sorted_docs

def get_chunks_by_ids(ids: list[int], conn) -> dict[int, str]:
    """Retrieve the content of chunks given their IDs."""
    if not ids:
        return {}
    with conn.cursor() as cur:
        cur.execute("SELECT id, content FROM document_chunks WHERE id = ANY(%s)", (ids,))
        rows = cur.fetchall()
        return {row[0]: row[1] for row in rows}

# --- Main Execution Logic ---
if __name__ == "__main__":
    model = initialize_model()
    try:
        conn = get_db_connection()

        # --- TEST CASE 1: Semantic Query ---
        print("\n--- Testing Semantic Query: 'What caused the budget overrun?' ---")
        query1 = "What caused the budget overrun?"
        fused_results1 = hybrid_search(query1, top_k=5, model=model, conn=conn)
        print("Fused RRF Scores:", fused_results1)
        top_doc_ids1 = [doc_id for doc_id, score in fused_results1[:3]]
        retrieved_content1 = get_chunks_by_ids(top_doc_ids1, conn)
        print("Top Retrieved Content:", retrieved_content1)

        # --- TEST CASE 2: Keyword-Specific Query ---
        print("\n--- Testing Keyword Query: 'What was the fix for INC-98765?' ---")
        query2 = "What was the fix for INC-98765?"
        fused_results2 = hybrid_search(query2, top_k=5, model=model, conn=conn)
        print("Fused RRF Scores:", fused_results2)
        top_doc_ids2 = [doc_id for doc_id, score in fused_results2[:3]]
        retrieved_content2 = get_chunks_by_ids(top_doc_ids2, conn)
        print("Top Retrieved Content:", retrieved_content2)

    except Exception as e:
        print(f"An error occurred: {e}")
    finally:
        if 'conn' in locals() and conn:
            conn.close()

# Before generating the prompt with the retrieved content
def reorder_chunks_for_llm(chunk_ids: list[int], rrf_scores: dict[int, float]) -> list[int]:
    sorted_chunks = sorted(chunk_ids, key=lambda id: rrf_scores.get(id, 0.0), reverse=True)
    reordered = []
    for i, chunk_id in enumerate(sorted_chunks):
        if i % 2 == 0:
            reordered.insert(0, chunk_id) # Add to the beginning
        else:
            reordered.append(chunk_id) # Add to the end
    return reordered

# Usage:
# top_doc_ids = [doc_id for doc_id, score in fused_results[:5]]
# reordered_ids = reorder_chunks_for_llm(top_doc_ids, dict(fused_results))
# retrieved_content = get_chunks_by_ids(reordered_ids, conn)
# # Now assemble the context string based on the new order

-- Set maintenance_work_mem to a higher value before building the index
SET maintenance_work_mem = '2GB';

-- Build the index after data has been loaded
CREATE INDEX ON document_chunks USING HNSW (embedding vector_cosine_ops);