Semantic Search with pgvector and Ollama

This solves semantic search when PostgreSQL is already your system of record. Store text and embeddings together, then use pgvector operators and indexes for nearest-neighbor lookup.

Code

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE IF NOT EXISTS text_embeddings (
    id SERIAL PRIMARY KEY,
    source TEXT,
    text TEXT NOT NULL,
    embedding vector(1024)
);

import os

import ollama
import psycopg2
from dotenv import load_dotenv

load_dotenv()


def get_connection():
    return psycopg2.connect(
        dbname=os.getenv("DB_NAME"),
        user=os.getenv("DB_USER"),
        password=os.getenv("DB_PASSWORD"),
        host=os.getenv("DB_HOST", "localhost"),
        port=os.getenv("DB_PORT", "5432"),
    )


def get_embedding(text, model="mxbai-embed-large"):
    response = ollama.embeddings(model=model, prompt=text)
    return response["embedding"]


def store_texts(rows):
    conn = get_connection()
    cursor = conn.cursor()
    for source, text in rows:
        embedding = get_embedding(text)
        cursor.execute(
            "INSERT INTO text_embeddings (source, text, embedding) VALUES (%s, %s, %s)",
            (source, text, embedding),
        )
    conn.commit()
    conn.close()


def find_similar(query_text, top_k=5):
    embedding = get_embedding(query_text)
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute(
        """
        SELECT source, text, 1 - (embedding <=> %s) AS similarity
        FROM text_embeddings
        ORDER BY embedding <=> %s
        LIMIT %s
        """,
        (embedding, embedding, top_k),
    )
    results = cursor.fetchall()
    conn.close()
    return results

Usage

store_texts([
    ("doc-1", "PostgreSQL can store embeddings with pgvector."),
    ("doc-2", "FAISS is a local vector search library."),
])

for source, text, similarity in find_similar("vector search in postgres"):
    print(f"{similarity:.3f}", source, text)

Indexes

Use HNSW when you want strong approximate nearest-neighbor search:

CREATE INDEX ON text_embeddings
USING hnsw (embedding vector_l2_ops);

SET hnsw.ef_search = 50;

Use IVF when you want a tunable lower-memory index:

CREATE INDEX ON text_embeddings
USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);

Requirements

Install PostgreSQL, pgvector, psycopg2, python-dotenv, and ollama. The vector dimension must match the embedding model output.

Source

The article points to example code in ernanhughes/pgvector-examples.

Full explanation

For installation notes, Windows setup details, and distance/index tradeoffs, read: PostgreSQL for AI: Storing and Searching Embeddings with pgvector.