Store embedding ke pgvector Postgres

Pinecone bagus, tapi $70/bulan untuk index kecil itu mahal. Kalau sudah pakai Postgres, pgvector cukup banget untuk sampai puluhan juta vector. Snippet ini setup table produk dengan embedding kolom + query top-K similar by cosine distance. Indonesia-ready.

Kode

import os
from dataclasses import dataclass
from typing import List

import psycopg
from openai import OpenAI

DATABASE_URL = os.environ["DATABASE_URL"]
openai_client = OpenAI()


@dataclass
class Produk:
    id: int
    nama: str
    deskripsi: str


def embed(text: str) -> list[float]:
    """Generate embedding 1536-dim dari OpenAI text-embedding-3-small."""
    resp = openai_client.embeddings.create(
        model="text-embedding-3-small",
        input=text,
    )
    return resp.data[0].embedding


def setup_schema(conn: psycopg.Connection) -> None:
    """Bikin extension pgvector + table produk."""
    with conn.cursor() as cur:
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS produk_embedding (
                id BIGSERIAL PRIMARY KEY,
                nama TEXT NOT NULL,
                deskripsi TEXT NOT NULL,
                embedding vector(1536) NOT NULL,
                created_at TIMESTAMPTZ DEFAULT now()
            )
            """
        )
        # Index HNSW untuk fast approximate similarity
        cur.execute(
            """
            CREATE INDEX IF NOT EXISTS produk_embedding_hnsw_idx
            ON produk_embedding USING hnsw (embedding vector_cosine_ops)
            WITH (m = 16, ef_construction = 64)
            """
        )
    conn.commit()


def insert_produk(conn: psycopg.Connection, nama: str, deskripsi: str) -> int:
    """Generate embedding + insert ke DB. Return ID baru."""
    teks_lengkap = f"{nama}. {deskripsi}"
    vec = embed(teks_lengkap)

    with conn.cursor() as cur:
        cur.execute(
            "INSERT INTO produk_embedding (nama, deskripsi, embedding) "
            "VALUES (%s, %s, %s) RETURNING id",
            (nama, deskripsi, vec),
        )
        new_id = cur.fetchone()[0]
    conn.commit()
    return new_id


def cari_mirip(
    conn: psycopg.Connection, query: str, top_k: int = 5
) -> List[tuple[Produk, float]]:
    """Cari produk paling mirip berdasarkan query bebas. Return (produk, distance)."""
    query_vec = embed(query)

    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT id, nama, deskripsi, embedding <=> %s::vector AS distance
            FROM produk_embedding
            ORDER BY embedding <=> %s::vector
            LIMIT %s
            """,
            (query_vec, query_vec, top_k),
        )
        rows = cur.fetchall()

    return [(Produk(id=r[0], nama=r[1], deskripsi=r[2]), float(r[3])) for r in rows]

Pemakaian

with psycopg.connect(DATABASE_URL) as conn:
    setup_schema(conn)

    # Insert beberapa produk Tokopedia
    insert_produk(conn, "Indomie Goreng Original", "Mie instan goreng rasa original, 85g per bungkus")
    insert_produk(conn, "Indomie Soto Lamongan", "Mie instan kuah rasa soto khas Lamongan")
    insert_produk(conn, "Aqua Galon 19L", "Air mineral kemasan galon 19 liter isi ulang")
    insert_produk(conn, "Lemonilo Mie Goreng", "Mie instan sehat tanpa MSG, varian goreng")

# Query semantic search
hasil = cari_mirip(conn, "mie instan rasa kuah nikmat", top_k=3)
for produk, dist in hasil:
    print(f"  [{dist:.3f}] {produk.nama}")
# Output:
#   [0.121] Indomie Soto Lamongan
#   [0.234] Indomie Goreng Original
#   [0.310] Lemonilo Mie Goreng

Kapan dipakai

• Semantic search produk e-commerce (Tokopedia, Bukalapak style).
• Deduplication artikel berita berdasarkan konten, bukan keyword.
• Recommendation: “produk mirip dengan yang baru dilihat”.
• RAG: chunk dokumen + retrieve relevant chunk untuk context LLM.

Catatan

• HNSW vs IVFFlat — HNSW lebih akurat tapi slower build. Untuk dataset di bawah 1 juta, HNSW worth it. Di atas itu pertimbangkan IVFFlat dengan lists = sqrt(N).
• Dimension matter — text-embedding-3-small = 1536, text-embedding-3-large = 3072. Storage 2x lebih besar, query latency 2x lebih lambat. Pilih yang sesuai kebutuhan.
• Distance operator: <=> cosine, <-> L2, <#> inner product. Cosine paling umum untuk text similarity.
• Hybrid search — kombinasikan dengan full-text search Postgres (tsvector) untuk dapet hasil terbaik. Vector miss kata-kata persis, BM25 miss makna.

Embedding deterministic-ish, tapi model upgrade (3-small → 3-large) bikin semua embedding lama harus di-regenerate. Track model version di kolom terpisah supaya tahu kapan harus reindex.