PostgreSQL Hybrid Search: Combining pg_vector and FTS for Advanced RAG

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

CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding VECTOR(1536) NOT NULL,
    content_tsv TSVECTOR,
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE OR REPLACE FUNCTION update_content_tsv()
RETURNS TRIGGER AS $$
BEGIN
    -- Using the 'english' configuration for stemming and stop words.
    -- Choose a configuration that matches your content language.
    NEW.content_tsv := to_tsvector('pg_catalog.english', NEW.content);
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

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

CREATE INDEX idx_documents_tsv ON documents USING GIN(content_tsv);

-- m: max number of connections per layer (16-64 is typical)
-- ef_construction: size of the dynamic candidate list during build (64-256 is typical)
CREATE INDEX idx_documents_embedding_hnsw ON documents USING HNSW (embedding vector_cosine_ops)
WITH (m = 32, ef_construction = 128);

-- Parameters for the query
-- :query_text - The user's raw text query (e.g., 'vulnerability in log4j')
-- :query_embedding - The embedding vector for the user's query
-- :match_limit - The number of results to fetch from each search (e.g., 20)
-- :final_limit - The final number of results to return (e.g., 10)

WITH fts_results AS (
    -- First CTE: Perform Full-Text Search
    SELECT
        id,
        -- ts_rank_cd normalizes for document length and is a good general-purpose choice
        ts_rank_cd(content_tsv, websearch_to_tsquery('english', :query_text)) AS fts_score,
        -- Assign a rank based on the score
        ROW_NUMBER() OVER (ORDER BY ts_rank_cd(content_tsv, websearch_to_tsquery('english', :query_text)) DESC) as fts_rank
    FROM documents
    -- The @@ operator performs the FTS match
    WHERE content_tsv @@ websearch_to_tsquery('english', :query_text)
    ORDER BY fts_score DESC
    LIMIT :match_limit
),
vector_results AS (
    -- Second CTE: Perform Vector Search
    SELECT
        id,
        -- The <=> operator gives cosine distance (0=identical, 2=opposite)
        -- We subtract from 1 to get similarity (1=identical, -1=opposite)
        1 - (embedding <=> :query_embedding) AS vector_score,
        ROW_NUMBER() OVER (ORDER BY embedding <=> :query_embedding ASC) as vector_rank
    FROM documents
    -- Exclude any exact matches already found by FTS to avoid trivial duplication in this CTE
    -- This is an optional optimization.
    WHERE id NOT IN (SELECT id FROM fts_results)
    ORDER BY embedding <=> :query_embedding ASC
    LIMIT :match_limit
),
-- Third CTE to union all results before fusion
all_results AS (
    SELECT id, fts_score, fts_rank, NULL::float8 as vector_score, NULL::bigint as vector_rank FROM fts_results
    UNION ALL
    SELECT id, NULL::float8, NULL::bigint, vector_score, vector_rank FROM vector_results
)
-- Final SELECT with RRF calculation
SELECT
    all_results.id,
    -- The RRF score calculation. k=60 is a common choice.
    SUM(1.0 / (60 + COALESCE(fts_rank, 2*:match_limit) + COALESCE(vector_rank, 2*:match_limit))) AS rrf_score,
    -- For debugging/inspection, you can also select the original ranks and scores
    MAX(fts_rank) as fts_rank,
    MAX(vector_rank) as vector_rank,
    MAX(fts_score) as fts_score,
    MAX(vector_score) as vector_score
FROM all_results
GROUP BY all_results.id
ORDER BY rrf_score DESC
LIMIT :final_limit;

EXPLAIN (ANALYZE, BUFFERS) <your_hybrid_query_here>;

Limit  (cost=...)
  ->  Sort  (cost=...)
        Sort Key: (sum(...)) DESC
        ->  HashAggregate  (cost=...)
              Group Key: all_results.id
              ->  Append  (cost=...)
                    ->  Subquery Scan on fts_results  (cost=...)
                          ->  Limit  (cost=...)
                                ->  Sort  (cost=...)
                                      Sort Key: (ts_rank_cd(...)) DESC
                                      ->  Bitmap Heap Scan on documents  (cost=...)
                                            Recheck Cond: (content_tsv @@ ...)
                                            ->  Bitmap Index Scan on idx_documents_tsv  (cost=...)
                                                  Index Cond: (content_tsv @@ ...)
                    ->  Subquery Scan on vector_results  (cost=...)
                          ->  Limit  (cost=...)
                                ->  Sort  (cost=...)
                                      Sort Key: ((embedding <=> ...)) ASC
                                      ->  Index Scan using idx_documents_embedding_hnsw on documents  (cost=...)
                                            Order By: (embedding <=> ...)

BEGIN;
SET LOCAL hnsw.ef_search = 100; -- Default is 40
-- Execute your hybrid query here
COMMIT;

import psycopg2
import numpy as np
from psycopg2.extras import RealDictCursor
from sentence_transformers import SentenceTransformer

# --- Configuration ---
DB_PARAMS = {
    'dbname': 'your_db',
    'user': 'your_user',
    'password': 'your_password',
    'host': 'localhost',
    'port': '5432'
}

# Use a pre-trained model for embeddings
# Ensure this matches the vector dimension in your DB (1536 for this one)
embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu') # Use 'cuda' for GPU

# --- Database Connection and Setup ---
def get_db_connection():
    # Register the vector type with psycopg2
    from pgvector.psycopg2 import register_vector
    conn = psycopg2.connect(**DB_PARAMS)
    register_vector(conn)
    return conn

# --- Hybrid Search Logic ---
HYBRID_SEARCH_SQL = """
WITH fts_results AS (
    SELECT
        id,
        ts_rank_cd(content_tsv, websearch_to_tsquery('english', %(query_text)s)) AS fts_score,
        ROW_NUMBER() OVER (ORDER BY ts_rank_cd(content_tsv, websearch_to_tsquery('english', %(query_text)s)) DESC) as fts_rank
    FROM documents
    WHERE content_tsv @@ websearch_to_tsquery('english', %(query_text)s)
    ORDER BY fts_score DESC
    LIMIT %(match_limit)s
),
vector_results AS (
    SELECT
        id,
        1 - (embedding <=> %(query_embedding)s) AS vector_score,
        ROW_NUMBER() OVER (ORDER BY embedding <=> %(query_embedding)s ASC) as vector_rank
    FROM documents
    ORDER BY embedding <=> %(query_embedding)s ASC
    LIMIT %(match_limit)s
),
all_results AS (
    SELECT id, fts_rank, NULL::bigint as vector_rank FROM fts_results
    UNION ALL
    SELECT id, NULL::bigint, vector_rank FROM vector_results
)
SELECT
    all_results.id,
    SUM(1.0 / (60 + COALESCE(fts_rank, 2*%(match_limit)s) + COALESCE(vector_rank, 2*%(match_limit)s))) AS rrf_score
FROM all_results
JOIN documents ON documents.id = all_results.id
GROUP BY all_results.id, documents.content
ORDER BY rrf_score DESC
LIMIT %(final_limit)s;
"""

def hybrid_search(query_text: str, match_limit: int = 20, final_limit: int = 10):
    conn = get_db_connection()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            # 1. Generate embedding for the query
            query_embedding = embedding_model.encode(query_text).tolist()
            
            params = {
                'query_text': query_text,
                'query_embedding': query_embedding,
                'match_limit': match_limit,
                'final_limit': final_limit
            }
            
            # 2. Set HNSW search parameter for this transaction
            cur.execute("SET LOCAL hnsw.ef_search = 100;")
            
            # 3. Execute the main query
            cur.execute(HYBRID_SEARCH_SQL, params)
            results = cur.fetchall()
            return results
    finally:
        conn.close()

# --- Example Usage ---
if __name__ == '__main__':
    # Assume you have populated your 'documents' table with some data
    # For example, one document about 'Log4j vulnerability CVE-2021-44228'
    # And another about 'Best practices for securing Java applications'
    
    # This query has a strong keyword ('log4j') and a semantic concept ('security best practices')
    search_query = "security best practices for log4j"
    
    print(f"Executing hybrid search for: '{search_query}'\n")
    search_results = hybrid_search(search_query)
    
    if search_results:
        for i, doc in enumerate(search_results):
            print(f"Rank {i+1} (ID: {doc['id']}, RRF Score: {doc['rrf_score']:.4f})")
    else:
        print("No results found.")