Scalable RAG with pgvector and Cross-Encoder Re-ranking

-- Ensure the pgvector extension is installed
CREATE EXTENSION IF NOT EXISTS vector;

-- Main table for document chunks and their embeddings
CREATE TABLE document_chunks (
    id BIGSERIAL PRIMARY KEY,
    document_id UUID NOT NULL,          -- Foreign key to a parent documents table
    content TEXT NOT NULL,              -- The actual text chunk
    token_count INT NOT NULL,           -- Useful for context window management
    embedding VECTOR(384) NOT NULL,     -- Dimensions must match your bi-encoder model
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- Add metadata for multi-tenancy or filtering
ALTER TABLE document_chunks ADD COLUMN tenant_id UUID;
CREATE INDEX idx_document_chunks_tenant_id ON document_chunks(tenant_id);

import psycopg2
import psycopg2.extras
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os

# --- Configuration ---
DB_CONNECTION_STRING = os.getenv("DATABASE_URL")
MODEL_NAME = 'all-MiniLM-L6-v2' # 384 dimensions
BATCH_SIZE = 128

# --- Initialize Models and Splitter ---
print("Loading bi-encoder model...")
bi_encoder = SentenceTransformer(MODEL_NAME)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=50)

def ingest_document(document_text: str, document_id: str, tenant_id: str):
    """Chunks, embeds, and batch-inserts a document into PostgreSQL."""
    chunks = text_splitter.split_text(document_text)
    
    print(f"Document split into {len(chunks)} chunks. Embedding in batches...")
    
    all_embeddings = bi_encoder.encode(chunks, show_progress_bar=True, batch_size=BATCH_SIZE)
    
    records_to_insert = []
    for i, chunk in enumerate(chunks):
        records_to_insert.append({
            'document_id': document_id,
            'content': chunk,
            'token_count': len(chunk.split()), # A rough approximation
            'embedding': all_embeddings[i].tolist(),
            'tenant_id': tenant_id
        })

    print(f"Inserting {len(records_to_insert)} records into the database...")
    with psycopg2.connect(DB_CONNECTION_STRING) as conn:
        with conn.cursor() as cur:
            psycopg2.extras.execute_values(
                cur,
                """INSERT INTO document_chunks (document_id, content, token_count, embedding, tenant_id)
                   VALUES %s""",
                [(r['document_id'], r['content'], r['token_count'], r['embedding'], r['tenant_id']) for r in records_to_insert]
            )
    print("Ingestion complete.")

# Example Usage:
if __name__ == '__main__':
    # In a real app, this would come from a file, API, etc.
    sample_doc_text = "... your large document text here ..."
    sample_doc_id = "a1b2c3d4-e5f6-a7b8-c9d0-e1f2a3b4c5d6"
    sample_tenant_id = "t1a2b3c4-d5e6-f7a8-b9c0-d1e2f3a4b5c6"
    ingest_document(sample_doc_text, sample_doc_id, sample_tenant_id)

import numpy as np

def fetch_candidates(query_text: str, tenant_id: str, k: int = 100) -> list[dict]:
    """Fetches top-k candidate chunks from PostgreSQL using vector similarity."""
    query_embedding = bi_encoder.encode(query_text)
    
    candidates = []
    with psycopg2.connect(DB_CONNECTION_STRING) as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
            # Set HNSW probe parameter for the current session
            # ef_search is the query-time equivalent of ef_construction
            cur.execute(f"SET hnsw.ef_search = {k * 2}") # A good starting point
            
            cur.execute(
                """SELECT id, content, document_id
                   FROM document_chunks
                   WHERE tenant_id = %s
                   ORDER BY embedding <=> %s
                   LIMIT %s""",
                (tenant_id, query_embedding.tolist(), k)
            )
            candidates = [dict(row) for row in cur.fetchall()]
    return candidates

from transformers import CrossEncoder

# --- Load Cross-Encoder Model (do this once at application startup) ---
print("Loading cross-encoder model...")
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

def rerank_candidates(query_text: str, candidates: list[dict], n: int = 5) -> list[dict]:
    """Re-ranks a list of candidates using a cross-encoder model."""
    if not candidates:
        return []
    
    # Format pairs for the cross-encoder: [ (query, doc_content), ... ]
    pairs = [(query_text, candidate['content']) for candidate in candidates]
    
    # Predict scores in a batch for efficiency
    print(f"Scoring {len(pairs)} pairs with cross-encoder...")
    scores = cross_encoder.predict(pairs, show_progress_bar=False)
    
    # Combine candidates with their new scores
    for i in range(len(candidates)):
        candidates[i]['rerank_score'] = scores[i]
        
    # Sort candidates by the new score in descending order
    candidates.sort(key=lambda x: x['rerank_score'], reverse=True)
    
    # Return the top-n most relevant candidates
    return candidates[:n]

# If you have a GPU, move the model to it
import torch
if torch.cuda.is_available():
    print("Moving cross-encoder to GPU")
    cross_encoder.model.to('cuda')

-- Add a tsvector column for FTS
ALTER TABLE document_chunks ADD COLUMN content_tsv TSVECTOR;

-- Create a trigger to automatically update the tsvector on insert/update
CREATE OR REPLACE FUNCTION update_content_tsv()
RETURNS TRIGGER AS $$
BEGIN
    NEW.content_tsv := to_tsvector('english', NEW.content);
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER tsvector_update
BEFORE INSERT OR UPDATE ON document_chunks
FOR EACH ROW EXECUTE FUNCTION update_content_tsv();

-- Create a GIN index for fast FTS
CREATE INDEX idx_gin_content_tsv ON document_chunks USING GIN(content_tsv);

def hybrid_fetch_candidates(query_text: str, tenant_id: str, k: int = 100):
    # This is a simplified implementation of RRF
    query_embedding = bi_encoder.encode(query_text)
    
    with psycopg2.connect(DB_CONNECTION_STRING) as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
            # Vector search results
            cur.execute(
                """SELECT id, content, (embedding <=> %s) as distance
                   FROM document_chunks WHERE tenant_id = %s
                   ORDER BY distance LIMIT %s""",
                (query_embedding.tolist(), tenant_id, k)
            )
            vector_results = {row['id']: {'rank': i + 1, 'data': row} for i, row in enumerate(cur.fetchall())}
            
            # FTS results
            cur.execute(
                """SELECT id, content, ts_rank(content_tsv, websearch_to_tsquery('english', %s)) as rank
                   FROM document_chunks WHERE tenant_id = %s AND content_tsv @@ websearch_to_tsquery('english', %s)
                   ORDER BY rank DESC LIMIT %s""",
                (query_text, tenant_id, query_text, k)
            )
            fts_results = {row['id']: {'rank': i + 1, 'data': row} for i, row in enumerate(cur.fetchall())}

    # RRF fusion logic
    rrf_scores = {}
    for doc_id, res in vector_results.items():
        rrf_scores[doc_id] = 1 / (60 + res['rank'])
    
    for doc_id, res in fts_results.items():
        if doc_id not in rrf_scores:
            rrf_scores[doc_id] = 0
        rrf_scores[doc_id] += 1 / (60 + res['rank'])

    # Sort by RRF score
    sorted_doc_ids = sorted(rrf_scores.keys(), key=lambda x: rrf_scores[x], reverse=True)
    
    # De-duplicate and combine results
    final_candidates = []
    seen_ids = set()
    for doc_id in sorted_doc_ids:
        if doc_id not in seen_ids:
            # Prioritize data from vector results if available
            data = (vector_results.get(doc_id) or fts_results.get(doc_id))['data']
            final_candidates.append(dict(data))
            seen_ids.add(doc_id)
        if len(final_candidates) >= k:
            break

    return final_candidates

class RAGPipeline:
    def __init__(self, db_conn_str, bi_encoder_model, cross_encoder_model):
        self.db_conn_str = db_conn_str
        self.bi_encoder = SentenceTransformer(bi_encoder_model)
        self.cross_encoder = CrossEncoder(cross_encoder_model)
        # Optional: Move to GPU if available
        if torch.cuda.is_available():
            self.cross_encoder.model.to('cuda')

    def _fetch_candidates(self, query_embedding, tenant_id: str, k: int = 100):
        # Implementation from Section 3.1
        # ... returns list of candidate dicts
        pass

    def _rerank_candidates(self, query_text: str, candidates: list[dict], n: int = 5):
        # Implementation from Section 3.2
        # ... returns list of re-ranked dicts
        pass

    def query(self, query_text: str, tenant_id: str, candidate_k: int = 100, final_n: int = 5):
        """Executes the full two-stage RAG pipeline."""
        print(f"Executing query for tenant '{tenant_id}': '{query_text}'")
        
        try:
            # Stage 1: Candidate Generation
            query_embedding = self.bi_encoder.encode(query_text)
            candidates = self._fetch_candidates(query_embedding, tenant_id, k=candidate_k)
            
            if not candidates:
                print("No candidates found.")
                return []

            print(f"Retrieved {len(candidates)} candidates from vector store.")

            # Stage 2: Re-ranking
            reranked_results = self._rerank_candidates(query_text, candidates, n=final_n)
            print(f"Re-ranked to top {len(reranked_results)} results.")

            return reranked_results

        except psycopg2.Error as e:
            print(f"Database error: {e}")
            # Add proper logging and error handling
            raise
        except Exception as e:
            print(f"An unexpected error occurred: {e}")
            raise

# --- Example Usage ---
if __name__ == '__main__':
    pipeline = RAGPipeline(
        db_conn_str=os.getenv("DATABASE_URL"),
        bi_encoder_model='all-MiniLM-L6-v2',
        cross_encoder_model='cross-encoder/ms-marco-MiniLM-L-6-v2'
    )
    
    user_query = "What are the performance implications of HNSW indexing in pgvector?"
    user_tenant_id = "t1a2b3c4-d5e6-f7a8-b9c0-d1e2f3a4b5c6"
    
    final_context_docs = pipeline.query(user_query, user_tenant_id)
    
    print("\n--- Final Context for LLM ---")
    for doc in final_context_docs:
        print(f"Score: {doc['rerank_score']:.4f}, ID: {doc['id']}")
        print(f"Content: {doc['content'][:150]}...")
        print("---")

    # This final context would then be formatted and sent to an LLM.