Semantic Search with pgvector and Python: Build a Production Search Engine

TL;DR — Build semantic search with pgvector and Python: embed text, store in PostgreSQL, query with cosine distance. MantraIdeas: "Database scans stored vectors to identify nearest neighbors using cosine similarity. Retrieves relevant content even when terminology differs." RedGate: "Measure how close two texts are by computing distance between vectors. Embeddings from OpenAI, Cohere, or sentence transformers. pgvector adds vector type and SQL distance operators to PostgreSQL." Virtido: "Moving from prototype to production requires reliability, performance, maintainability. Compare OpenAI vs Cohere vs open-source on your data. Track metrics over time." Medium: "Semantic and hybrid search with Python and pgvector. Efficient similarity searches using cosine similarity or Euclidean distance." Learn more with pgvector tutorial, embeddings, embedding models, and cosine similarity.

MantraIdeas explains the core concept: "The database then scans its stored vectors to identify the nearest neighbors to the query vector, using distance metrics like cosine similarity to quantify how closely they align. This methodology allows us to retrieve relevant content even when the terminology used in the query and the stored data differs."

RedGate provides the foundation: "We can measure how close two pieces of text are by computing the distance between their vectors. AI-generated embeddings come from models trained on massive datasets. Models built by OpenAI or Cohere, or sentence transformers, convert text into these numeric representations. pgvector adds a vector data type and SQL distance operators to PostgreSQL."

Semantic Search Architecture

flowchart TD subgraph Indexing["Document Indexing"] Docs["Documents
files, DB, API"] Batch["Batch Embed
100-1000 at a time"] Model["Embedding Model
OpenAI / sentence-transformers / BGE"] Embeddings["Vector Embeddings
1536-d or 384-d"] Insert["Bulk Insert
executemany or COPY"] PG["PostgreSQL + pgvector
vector(N) column"] HNSW["HNSW Index
m=16, ef_construction=200"] end subgraph Search["Query Search"] Query["User Query
text string"] QueryEmbed["Embed Query
same model"] SQL["SQL Query
ORDER BY embedding <=> query
LIMIT k"] Filter["Metadata Filter
WHERE clause"] Results["Ranked Results
with similarity scores"] end subgraph API["FastAPI Endpoint"] Endpoint["/search endpoint
POST query + filters"] Pool["Connection Pool
psycopg2 pool"] Cache["Response Cache
frequent queries"] Response["JSON Response
results + metadata"] end Docs --> Batch --> Model --> Embeddings --> Insert --> PG PG --> HNSW Query --> QueryEmbed --> SQL HNSW --> SQL Filter --> SQL SQL --> Results Endpoint --> Pool --> SQL Cache --> Endpoint Results --> Response
Model Dimensions Cost Latency Quality Best For
OpenAI text-embedding-3-small 1536 $0.02/1M tokens Fast High Production, managed
OpenAI text-embedding-3-large 3072 $0.13/1M tokens Medium Highest Maximum quality
all-MiniLM-L6-v2 384 Free Very fast Good Prototyping, local
BGE-large-en-v1.5 1024 Free Medium High Best open-source
Cohere embed-english-v3 1024 Paid Fast High Multilingual
nomic-embed-text 768 Free Fast Good Self-hosted

Implementation

import os
import time
import json
from dataclasses import dataclass, field
from typing import Optional
from enum import Enum
import psycopg2
from pgvector.psycopg2 import register_vector

class EmbeddingModel(Enum):
    OPENAI_SMALL = "openai_small"
    OPENAI_LARGE = "openai_large"
    MINILM = "minilm"
    BGE_LARGE = "bge_large"
    COHERE = "cohere"
    NOMIC = "nomic"

@dataclass
class SemanticSearchEngine:
    """Production semantic search engine with pgvector and Python."""

    def __init__(self, dsn: str,
                 model: EmbeddingModel = EmbeddingModel.OPENAI_SMALL,
                 dimensions: int = 1536,
                 top_k: int = 5,
                 ef_search: int = 100):
        self.dsn = dsn
        self.model = model
        self.dimensions = dimensions
        self.top_k = top_k
        self.ef_search = ef_search
        self._conn = None
        self._embedder = None

    @property
    def conn(self):
        """Lazy connection with pgvector registered."""
        if self._conn is None:
            self._conn = psycopg2.connect(self.dsn)
            register_vector(self._conn)
        return self._conn

    @property
    def embedder(self):
        """Lazy embedder initialization."""
        if self._embedder is None:
            self._embedder = self._init_embedder()
        return self._embedder

    def _init_embedder(self):
        """Initialize embedding model."""
        if self.model == EmbeddingModel.MINILM:
            from sentence_transformers import SentenceTransformer
            return SentenceTransformer(
                'all-MiniLM-L6-v2')
        elif self.model == EmbeddingModel.BGE_LARGE:
            from sentence_transformers import SentenceTransformer
            return SentenceTransformer(
                'BAAI/bge-large-en-v1.5')
        elif self.model in (EmbeddingModel.OPENAI_SMALL,
                            EmbeddingModel.OPENAI_LARGE):
            from openai import OpenAI
            client = OpenAI()
            model_name = (
                "text-embedding-3-small"
                if self.model == EmbeddingModel.OPENAI_SMALL
                else "text-embedding-3-large"
            )
            return lambda text: client.embeddings.create(
                input=text, model=model_name
            ).data[0].embedding
        return None

    def embed_text(self, text: str) -> list:
        """Embed a single text."""
        if self.model in (EmbeddingModel.MINILM,
                          EmbeddingModel.BGE_LARGE):
            return self.embedder.encode(text).tolist()
        else:
            return self.embedder(text)

    def embed_batch(self, texts: list) -> list:
        """Embed a batch of texts."""
        if self.model in (EmbeddingModel.MINILM,
                          EmbeddingModel.BGE_LARGE):
            return self.embedder.encode(texts).tolist()
        else:
            from openai import OpenAI
            client = OpenAI()
            model_name = (
                "text-embedding-3-small"
                if self.model == EmbeddingModel.OPENAI_SMALL
                else "text-embedding-3-large"
            )
            response = client.embeddings.create(
                input=texts, model=model_name
            )
            return [d.embedding for d in response.data]

    def setup_database(self):
        """Create extension, table, and index."""
        cur = self.conn.cursor()
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
        cur.execute(f"""
            CREATE TABLE IF NOT EXISTS documents (
                id BIGSERIAL PRIMARY KEY,
                title TEXT NOT NULL,
                content TEXT NOT NULL,
                embedding vector({self.dimensions}) NOT NULL,
                source TEXT,
                category TEXT,
                created_at TIMESTAMP DEFAULT NOW()
            );
        """)
        cur.execute(f"""
            CREATE INDEX IF NOT EXISTS documents_embedding_hnsw_idx
            ON documents USING hnsw (embedding vector_cosine_ops)
            WITH (m = 16, ef_construction = 200);
        """)
        self.conn.commit()

    def index_documents(self, documents: list,
                        batch_size: int = 100) -> dict:
        """Batch index documents with embeddings."""
        cur = self.conn.cursor()
        total = 0

        for i in range(0, len(documents), batch_size):
            batch = documents[i:i + batch_size]
            texts = [doc["content"] for doc in batch]
            embeddings = self.embed_batch(texts)

            for doc, embedding in zip(batch, embeddings):
                cur.execute(
                    """
                    INSERT INTO documents
                        (title, content, embedding, source, category)
                    VALUES (%s, %s, %s, %s, %s)
                    ON CONFLICT (id) DO UPDATE SET
                        title = EXCLUDED.title,
                        content = EXCLUDED.content,
                        embedding = EXCLUDED.embedding,
                        source = EXCLUDED.source,
                        category = EXCLUDED.category
                    """,
                    (doc["title"], doc["content"],
                     embedding, doc.get("source", ""),
                     doc.get("category", "")),
                )
            total += len(batch)

        self.conn.commit()
        return {
            "indexed": total,
            "batches": (len(documents) + batch_size - 1) // batch_size,
            "model": self.model.value,
            "dimensions": self.dimensions,
        }

    def search(self, query: str,
               top_k: int = None,
               category: str = None,
               min_similarity: float = 0.0) -> dict:
        """Semantic search with optional filtering."""
        k = top_k or self.top_k
        query_embedding = self.embed_text(query)

        cur = self.conn.cursor()

        with self.conn:
            cur.execute(
                "SET LOCAL hnsw.ef_search = %s",
                (self.ef_search,))

            if category:
                cur.execute(
                    """
                    SELECT id, title, content, source, category,
                           1 - (embedding <=> %s::vector) AS similarity
                    FROM documents
                    WHERE category = %s
                    ORDER BY embedding <=> %s::vector
                    LIMIT %s
                    """,
                    (query_embedding, category,
                     query_embedding, k),
                )
            else:
                cur.execute(
                    """
                    SELECT id, title, content, source, category,
                           1 - (embedding <=> %s::vector) AS similarity
                    FROM documents
                    ORDER BY embedding <=> %s::vector
                    LIMIT %s
                    """,
                    (query_embedding, query_embedding, k),
                )

            results = cur.fetchall()

        filtered = [
            {
                "id": r[0],
                "title": r[1],
                "content": r[2],
                "source": r[3],
                "category": r[4],
                "similarity": float(r[5]),
            }
            for r in results
            if float(r[5]) >= min_similarity
        ]

        return {
            "query": query,
            "results": filtered,
            "count": len(filtered),
            "model": self.model.value,
            "ef_search": self.ef_search,
        }

    def hybrid_search(self, query: str,
                      top_k: int = None,
                      alpha: float = 0.5) -> dict:
        """Hybrid search: semantic + keyword using RRF fusion.

        alpha: weight for semantic (1-alpha for keyword)
        Uses PostgreSQL tsvector for keyword search.
        """
        k = top_k or self.top_k
        query_embedding = self.embed_text(query)

        cur = self.conn.cursor()

        with self.conn:
            cur.execute(
                "SET LOCAL hnsw.ef_search = %s",
                (self.ef_search,))

            # Semantic search
            cur.execute(
                """
                SELECT id, title, content, source,
                       ROW_NUMBER() OVER (
                           ORDER BY embedding <=> %s::vector
                       ) AS vector_rank
                FROM documents
                ORDER BY embedding <=> %s::vector
                LIMIT %s
                """,
                (query_embedding, query_embedding, k * 2),
            )
            vector_results = {r[0]: r for r in cur.fetchall()}

            # Keyword search using tsvector
            cur.execute(
                """
                SELECT id, title, content, source,
                       ROW_NUMBER() OVER (
                           ORDER BY ts_rank(
                               to_tsvector('english', content),
                               plainto_tsquery('english', %s)
                           ) DESC
                       ) AS keyword_rank
                FROM documents
                WHERE to_tsvector('english', content)
                      @@ plainto_tsquery('english', %s)
                LIMIT %s
                """,
                (query, query, k * 2),
            )
            keyword_results = {r[0]: r for r in cur.fetchall()}

            # Reciprocal Rank Fusion
            rrf_k = 60
            all_ids = set(vector_results.keys()) | \
                      set(keyword_results.keys())
            fused = []

            for doc_id in all_ids:
                v_rank = vector_results.get(doc_id)
                k_rank = keyword_results.get(doc_id)

                v_score = (1 / (rrf_k + v_rank[4])
                           if v_rank else 0)
                k_score = (1 / (rrf_k + k_rank[4])
                           if k_rank else 0)

                combined = alpha * v_score + (1 - alpha) * k_score

                row = v_rank or k_rank
                fused.append({
                    "id": doc_id,
                    "title": row[1],
                    "content": row[2],
                    "source": row[3],
                    "score": combined,
                    "vector_rank": v_rank[4] if v_rank else None,
                    "keyword_rank": k_rank[4] if k_rank else None,
                })

            fused.sort(key=lambda x: x["score"], reverse=True)

        return {
            "query": query,
            "results": fused[:k],
            "count": len(fused[:k]),
            "alpha": alpha,
            "fusion": "RRF",
        }

FastAPI Search Endpoint

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional

app = FastAPI(title="Semantic Search API")
engine = SemanticSearchEngine(
    dsn="postgresql://user:pass@localhost/db",
    model=EmbeddingModel.OPENAI_SMALL,
)

class SearchRequest(BaseModel):
    query: str
    top_k: Optional[int] = 5
    category: Optional[str] = None
    min_similarity: Optional[float] = 0.0

class SearchResult(BaseModel):
    id: int
    title: str
    content: str
    source: str
    category: str
    similarity: float

class SearchResponse(BaseModel):
    query: str
    results: list[SearchResult]
    count: int

@app.post("/search", response_model=SearchResponse)
async def search(req: SearchRequest):
    results = engine.search(
        query=req.query,
        top_k=req.top_k,
        category=req.category,
        min_similarity=req.min_similarity,
    )
    return SearchResponse(
        query=results["query"],
        results=[SearchResult(**r) for r in results["results"]],
        count=results["count"],
    )

@app.post("/hybrid-search")
async def hybrid_search(req: SearchRequest):
    results = engine.hybrid_search(
        query=req.query,
        top_k=req.top_k,
    )
    return results

Semantic Search Checklist

  • [ ] Install pgvector: CREATE EXTENSION IF NOT EXISTS vector;
  • [ ] Install Python libs: pip install psycopg2-binary pgvector sentence-transformers
  • [ ] For OpenAI: pip install openai
  • [ ] Create table with vector(N) column matching embedding model dimensions
  • [ ] Build HNSW index: CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 200)
  • [ ] Choose embedding model: OpenAI for production, sentence-transformers for local/free
  • [ ] OpenAI text-embedding-3-small: 1536 dims, $0.02/1M tokens, fast, high quality
  • [ ] all-MiniLM-L6-v2: 384 dims, free, very fast, good for prototyping
  • [ ] BGE-large-en-v1.5: 1024 dims, free, best open-source quality
  • [ ] Use same embedding model for indexing and querying — mismatched models produce garbage
  • [ ] Batch embed documents: 100-1000 at a time for efficiency
  • [ ] Use psycopg2 executemany or COPY for fast bulk inserts
  • [ ] Build HNSW index after all documents are inserted (faster than incremental)
  • [ ] Use ON CONFLICT DO UPDATE for idempotent re-indexing
  • [ ] Set maintenance_work_mem = '2GB' before building index
  • [ ] Semantic search: embed query, then SELECT ... ORDER BY embedding <=> query_vector LIMIT k
  • [ ] Use <=> cosine distance operator for text/semantic embeddings
  • [ ] Use 1 - (embedding <=> query) to convert distance to similarity score (0-1)
  • [ ] Apply metadata filtering with WHERE clauses: WHERE category = 'tech'
  • [ ] Use SET LOCAL hnsw.ef_search = 100 in transaction for tuning
  • [ ] Create FastAPI /search endpoint: accept query, embed, search, return JSON
  • [ ] Use connection pooling for production (psycopg2 pool or pgbouncer)
  • [ ] Add response caching for frequent queries
  • [ ] Add rate limiting to search endpoint
  • [ ] Track search metrics: latency, recall, click-through rate
  • [ ] Hybrid search: combine semantic + keyword using RRF fusion
  • [ ] RRF formula: score = alpha * 1/(k + rank_vector) + (1-alpha) * 1/(k + rank_keyword), k=60
  • [ ] Use ParadeDB for native BM25 + vector hybrid inside Postgres
  • [ ] Manual hybrid: run vector search and tsvector search separately, fuse with RRF
  • [ ] Semantic search retrieves relevant content even when terminology differs
  • [ ] AI-generated embeddings from models trained on massive datasets
  • [ ] pgvector adds vector data type and SQL distance operators to PostgreSQL
  • [ ] Compare OpenAI vs Cohere vs open-source on your data before choosing
  • [ ] Track metrics over time as content and query patterns evolve
  • [ ] Moving from prototype to production requires reliability, performance, maintainability
  • [ ] Monitor index size: SELECT pg_size_pretty(pg_relation_size('documents_embedding_hnsw_idx'))
  • [ ] Monitor query latency: EXPLAIN ANALYZE on search queries
  • [ ] Monitor recall: compare ANN results with exact search on sample
  • [ ] Read pgvector tutorial for setup
  • [ ] Read embeddings for fundamentals
  • [ ] Read embedding models for selection
  • [ ] Read cosine similarity for metrics
  • [ ] Read hybrid search for BM25 + vector
  • [ ] Read what is RAG for RAG architecture
  • [ ] Read vector databases for ANN fundamentals
  • [ ] Test: semantic search finds relevant docs despite different terminology
  • [ ] Test: batch indexing inserts all documents correctly
  • [ ] Test: FastAPI endpoint returns ranked results with similarity scores
  • [ ] Test: metadata filtering narrows results correctly
  • [ ] Test: hybrid search combines semantic and keyword results with RRF
  • [ ] Test: same embedding model used for indexing and querying
  • [ ] Document model choice, dimensions, index parameters, and API endpoints

FAQ

How do you build a semantic search engine with pgvector and Python?

Build semantic search by embedding text with sentence-transformers or OpenAI, storing vectors in PostgreSQL with pgvector, and querying with cosine distance. MantraIdeas: "The database scans stored vectors to identify nearest neighbors to the query vector, using distance metrics like cosine similarity. This allows retrieving relevant content even when the terminology used in the query and stored data differs." RedGate: "We can measure how close two pieces of text are by computing the distance between their vectors. AI-generated embeddings from OpenAI, Cohere, or sentence transformers convert text into numeric representations. pgvector adds vector data type and SQL distance operators to PostgreSQL." Medium: "Semantic search and hybrid search using Python with pgvector. Supports efficient similarity searches using metrics like cosine similarity or Euclidean distance." Steps: (1) Install pgvector and sentence-transformers. (2) Create table with vector column. (3) Embed documents and insert. (4) Build HNSW index. (5) Embed query and search with ORDER BY embedding <=> query_vector LIMIT k.

What embedding models work best for semantic search with pgvector?

Use OpenAI text-embedding-3-small for production, sentence-transformers all-MiniLM-L6-v2 for local/free, and BGE-large-en-v1.5 for best open-source quality. RedGate: "Models built by OpenAI or Cohere, or sentence transformers, convert text into numeric representations. AI-generated embeddings come from models trained on massive datasets." Virtido: "Embedding model — Compare OpenAI vs Cohere vs open-source on your data. Track metrics over time as your content and query patterns evolve." Model comparison: (1) OpenAI text-embedding-3-small: 1536 dims, fast, $0.02/1M tokens, best for production. (2) sentence-transformers all-MiniLM-L6-v2: 384 dims, free, local, fast, good for prototyping. (3) BGE-large-en-v1.5: 1024 dims, free, best open-source quality. (4) Cohere embed-english-v3: 1024 dims, good multilingual. Choose based on cost, latency, quality, and whether you need self-hosted.

How do you implement a FastAPI search endpoint with pgvector?

Create a FastAPI endpoint that embeds the query, searches pgvector with cosine distance, and returns ranked results with metadata. Virtido: "Moving from prototype to production requires attention to reliability, performance, and maintainability. from fastapi import FastAPI, HTTPException from pydantic." Implementation: (1) Create FastAPI app with /search endpoint. (2) Accept query string and optional filters. (3) Embed query with same model used for indexing. (4) Execute SQL: SELECT id, title, content, embedding <=> query_vector AS distance FROM documents ORDER BY embedding <=> query_vector LIMIT k. (5) Return results with similarity scores. (6) Add metadata filtering with WHERE clauses. (7) Use connection pooling for production. (8) Add rate limiting and caching. (9) Use SET LOCAL hnsw.ef_search for tuning.

How do you do batch indexing for semantic search with pgvector?

Batch index documents by embedding in batches, using executemany or COPY for fast inserts, and building the HNSW index after insertion. MantraIdeas: "The database scans stored vectors to identify nearest neighbors using distance metrics like cosine similarity." RedGate: "pgvector adds vector data type and SQL distance operators to PostgreSQL." Batch indexing: (1) Load documents from source (files, database, API). (2) Embed in batches of 100-1000 using sentence-transformers or OpenAI API. (3) Use psycopg2 executemany or COPY for fast bulk insert. (4) Build HNSW index after all documents are inserted (faster than incremental). (5) Use ON CONFLICT DO UPDATE for idempotent re-indexing. (6) Track indexing progress and errors. (7) Set maintenance_work_mem = '2GB' before building index. (8) Monitor index build time and size.

How do you combine semantic search with keyword search in pgvector?

Combine semantic (vector) search with keyword (BM25) search using ParadeDB or manual fusion in SQL for hybrid search. Medium: "Semantic search and hybrid search using Python with pgvector. Supports efficient similarity searches using metrics like cosine similarity or Euclidean distance." Virtido: "Track metrics over time as your content and query patterns evolve." Hybrid search approaches: (1) ParadeDB: install pg_search extension for BM25 inside Postgres, combine with pgvector using RRF fusion. (2) Manual fusion: run vector search and tsvector full-text search separately, combine scores with weighted average or RRF. (3) SQL UNION: SELECT from both searches, order by combined score. (4) Reciprocal Rank Fusion (RRF): score = 1/(k + rank_vector) + 1/(k + rank_keyword), where k=60. ParadeDB is the cleanest approach — BM25 and vector search in one SQL query inside Postgres.


Want a self-hosted AI company brain that does all of this out of the box?
Book a demo →