Production RAG: Hybrid Search with PostgreSQL, pgvector & BM25

# Simplified demonstration of the problem
from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer('all-MiniLM-L6-v2')

chunks = [
    "The investigation into the 'Heimdall' security incident (ref: HMDL-2024-01) revealed a buffer overflow vulnerability. The patch was released in Q3.",
    "A critical security flaw was discovered in our primary authentication service, leading to a widespread system compromise. This event highlighted the need for better memory safety protocols.",
    "Quarterly budget report for the platform engineering team shows increased spending on monitoring tools."
]

query = "details on incident HMDL-2024-01"

# Encode chunks and query
chunk_embeddings = model.encode(chunks, convert_to_tensor=True)
query_embedding = model.encode(query, convert_to_tensor=True)

# Compute cosine similarities
cosine_scores = util.pytorch_cos_sim(query_embedding, chunk_embeddings)[0]

print("Query:", query)
for i, score in enumerate(cosine_scores):
    print(f"Score: {score:.4f} - Chunk {i+1}: {chunks[i][:30]}...")

# In many runs, the scores can be surprisingly close, or even wrong.
# A potential output might be:
# Query: details on incident HMDL-2024-01
# Score: 0.6812 - Chunk 1: The investigation into the 'He...
# Score: 0.6590 - Chunk 2: A critical security flaw was d...
# Score: 0.1234 - Chunk 3: Quarterly budget report for t...

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

CREATE TABLE document_chunks (
    id SERIAL PRIMARY KEY,
    source_id VARCHAR(255) NOT NULL, -- Identifier for the original document
    chunk_sequence INT NOT NULL,      -- Order of the chunk within the document
    content TEXT NOT NULL,
    embedding VECTOR(384) NOT NULL, -- Assuming all-MiniLM-L6-v2 model (384 dims)
    content_tsv TSVECTOR,
    created_at TIMESTAMPTZ DEFAULT NOW(),
    UNIQUE (source_id, chunk_sequence)
);

pip install sentence-transformers langchain psycopg2-binary

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

# --- Configuration ---
DB_CONNECTION_STRING = "postgresql://user:password@host:port/dbname"
MODEL_NAME = 'all-MiniLM-L6-v2'

# --- Initialize components ---
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    length_function=len,
)
model = SentenceTransformer(MODEL_NAME)

def ingest_document(source_id: str, document_text: str):
    """Chunks, embeds, and stores a single document."""
    print(f"Ingesting document: {source_id}")
    
    # 1. Chunk the document
    chunks = text_splitter.split_text(document_text)
    
    # 2. Generate dense embeddings
    embeddings = model.encode(chunks, show_progress_bar=True)
    
    # 3. Prepare data for insertion
    records_to_insert = []
    for i, chunk in enumerate(chunks):
        embedding = embeddings[i].tolist()
        # The content_tsv will be generated by a trigger in the DB for efficiency
        records_to_insert.append((source_id, i, chunk, embedding))

    # 4. Insert into database
    with psycopg2.connect(DB_CONNECTION_STRING) as conn:
        with conn.cursor() as cur:
            execute_values(
                cur,
                """INSERT INTO document_chunks (source_id, chunk_sequence, content, embedding)
                   VALUES %s""",
                records_to_insert
            )
    print(f"Successfully ingested {len(chunks)} chunks for document: {source_id}")

-- Create a function to update the tsvector column
CREATE OR REPLACE FUNCTION update_document_chunks_tsv() RETURNS TRIGGER AS $$
BEGIN
    NEW.content_tsv := to_tsvector('english', NEW.content);
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- Create a trigger that calls the function before insert or update
CREATE TRIGGER tsvector_update_trigger
BEFORE INSERT OR UPDATE ON document_chunks
FOR EACH ROW EXECUTE FUNCTION update_document_chunks_tsv();

    CREATE INDEX idx_gin_document_chunks_tsv ON document_chunks USING GIN(content_tsv);

    -- HNSW index for vector search (cosine distance)
    -- The choice of m and ef_construction are critical tuning parameters.
    -- m: max number of connections per layer (16-64 is typical)
    -- ef_construction: size of the dynamic candidate list during build (higher is better recall but slower build)
    CREATE INDEX idx_hnsw_document_chunks_embedding ON document_chunks USING HNSW (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 64);

import psycopg2
from psycopg2.extras import DictCursor
from sentence_transformers import SentenceTransformer

# --- Re-use configuration from ingestion ---
DB_CONNECTION_STRING = "postgresql://user:password@host:port/dbname"
MODEL_NAME = 'all-MiniLM-L6-v2'
model = SentenceTransformer(MODEL_NAME)

def reciprocal_rank_fusion(results: list[list[dict]], k: int = 60) -> list[dict]:
    """Performs RRF on a list of search result lists."""
    fused_scores = {}

    # results is a list of lists, e.g., [[doc1, doc2], [doc2, doc3]]
    for result_set in results:
        for rank, doc in enumerate(result_set):
            doc_id = doc['id']
            if doc_id not in fused_scores:
                fused_scores[doc_id] = {'score': 0, 'doc': doc}
            
            fused_scores[doc_id]['score'] += 1 / (k + rank)

    # Sort by fused score in descending order
    reranked_results = sorted(fused_scores.values(), key=lambda x: x['score'], reverse=True)
    
    return [item['doc'] for item in reranked_results]

def hybrid_search(query: str, top_k: int = 10) -> list[dict]:
    """Performs hybrid search and returns re-ranked documents."""
    
    # 1. Generate query vectors
    query_embedding = model.encode(query).tolist()
    
    # For FTS, we convert the query to a tsquery format, using '&' for AND logic
    # This is a simple approach; more sophisticated logic can be used
    query_fts = " & ".join(query.split())
    
    vector_results = []
    fts_results = []

    with psycopg2.connect(DB_CONNECTION_STRING, cursor_factory=DictCursor) as conn:
        with conn.cursor() as cur:
            # 2. Perform Vector Search
            # <=> is the cosine distance operator
            cur.execute(
                """SELECT id, source_id, content, (1 - (embedding <=> %s)) as similarity
                   FROM document_chunks 
                   ORDER BY embedding <=> %s 
                   LIMIT %s""",
                (str(query_embedding), str(query_embedding), top_k * 2) # Fetch more to allow for fusion
            )
            vector_results = cur.fetchall()

            # 3. Perform Full-Text Search
            cur.execute(
                """SELECT id, source_id, content, ts_rank(content_tsv, to_tsquery('english', %s)) as rank
                   FROM document_chunks 
                   WHERE content_tsv @@ to_tsquery('english', %s) 
                   ORDER BY rank DESC 
                   LIMIT %s""",
                (query_fts, query_fts, top_k * 2)
            )
            fts_results = cur.fetchall()

    # 4. Fuse the results using RRF
    # Convert psycopg2's DictRow objects to standard dicts
    vector_results_dicts = [dict(row) for row in vector_results]
    fts_results_dicts = [dict(row) for row in fts_results]
    
    reranked_results = reciprocal_rank_fusion([vector_results_dicts, fts_results_dicts])
    
    return reranked_results[:top_k]

# --- Example Usage ---
if __name__ == '__main__':
    # Assume some documents have been ingested
    # e.g., ingest_document("security_policy_v1", "The policy document HMDL-2024-01 outlines our response...")
    
    search_query = "details on incident HMDL-2024-01"
    results = hybrid_search(search_query)
    
    print(f"Hybrid search results for: '{search_query}'\n")
    for res in results:
        print(f"ID: {res['id']}, Source: {res['source_id']}")
        print(f"Content: {res['content'][:200]}...\n")

-- Example of setting it for a session
SET hnsw.ef_search = 100;
SELECT id FROM document_chunks ORDER BY embedding <=> '[...]' LIMIT 10;

from sentence_transformers import CrossEncoder

# This model is specifically trained for re-ranking
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

def rerank_with_cross_encoder(query: str, documents: list[dict]) -> list[dict]:
    """Re-ranks a list of documents using a CrossEncoder model."""
    # The model expects a list of [query, document_content] pairs
    pairs = [[query, doc['content']] for doc in documents]
    
    # Get scores from the model
    scores = cross_encoder.predict(pairs, show_progress_bar=False)
    
    # Add scores to documents and sort
    for i, doc in enumerate(documents):
        doc['cross_encoder_score'] = scores[i]
        
    return sorted(documents, key=lambda x: x['cross_encoder_score'], reverse=True)

# --- In your main workflow ---
# 1. Get initial candidates from hybrid search
candidates = hybrid_search(search_query, top_k=50)

# 2. Re-rank the top candidates for maximum precision
final_results = rerank_with_cross_encoder(search_query, candidates)

# 3. Use the top 5-10 of these for the LLM context
top_5_final = final_results[:5]

import openai

# Assume openai.api_key is configured

def generate_answer(query: str, context_chunks: list[dict]):
    """Generates an answer using the LLM based on the provided context."""
    
    context_str = "\n\n---\n\n".join([chunk['content'] for chunk in context_chunks])
    
    prompt = f"""
    You are an expert Q&A system. Your answer must be grounded in the provided context.
    Do not use any information outside of the context. If the context does not contain
    the answer, state that you cannot answer the question.

    CONTEXT:
    {context_str}

    QUESTION:
    {query}

    ANSWER:
    """

    response = openai.chat.completions.create(
        model="gpt-4-turbo-preview",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.1,
    )
    
    return response.choices[0].message.content

# --- Final End-to-End Example ---
search_query = "What was the resolution for incident HMDL-2024-01?"

# 1. Hybrid search to get candidates
candidates = hybrid_search(search_query, top_k=20)

# 2. (Optional but recommended) Cross-encoder re-ranking
reranked = rerank_with_cross_encoder(search_query, candidates)

# 3. Select top N for context
final_context = reranked[:5]

# 4. Generate final answer
answer = generate_answer(search_query, final_context)

print("--- Generated Answer ---")
print(answer)