SQLAlchemy + Postgres Cheatsheet 12 — pgvector (Embeddings)

CREATE EXTENSION vector;

uv add pgvector

from pgvector.sqlalchemy import Vector

class Document(Base):
    __tablename__ = "documents"
    id: Mapped[int] = mapped_column(primary_key=True)
    body: Mapped[str]
    embedding: Mapped[list[float]] = mapped_column(Vector(1024))

__table_args__ = (
    Index(
        "ix_docs_embedding_hnsw",
        "embedding",
        postgresql_using="hnsw",
        postgresql_ops={"embedding": "vector_cosine_ops"},
    ),
)

__table_args__ = (
    Index(
        "ix_docs_embedding_ivf",
        "embedding",
        postgresql_using="ivfflat",
        postgresql_ops={"embedding": "vector_cosine_ops"},
        postgresql_with={"lists": 100},   # ~sqrt(rows)
    ),
)

q_emb = await embed("query text")

stmt = (
    select(Document, Document.embedding.cosine_distance(q_emb).label("dist"))
    .order_by("dist")
    .limit(10)
)

from sqlalchemy import literal_column

stmt = (
    select(Document)
    .order_by(Document.embedding.op("<=>")(q_emb))
    .limit(10)
)

stmt = (
    select(Document)
    .where(Document.user_id == user_id)
    .order_by(Document.embedding.cosine_distance(q_emb))
    .limit(10)
)

async with session.begin():
    await session.execute(text("SET LOCAL hnsw.ef_search = 200"))
    result = await session.execute(stmt)

postgresql_with={"m": 16, "ef_construction": 64}

similarity_col = (1 - Document.embedding.cosine_distance(q_emb)).label("similarity")
stmt = select(Document, similarity_col).order_by(similarity_col.desc()).limit(10)

stmt = (
    select(Document)
    .where(Document.embedding.cosine_distance(q_emb) < 0.5)
    .order_by(Document.embedding.cosine_distance(q_emb))
    .limit(10)
)

records = [
    {"body": doc, "embedding": await embed(doc)}
    for doc in docs
]
await session.execute(insert(Document), records)
await session.commit()

from pgvector.sqlalchemy import HALFVEC

class Document(Base):
    embedding: Mapped[list[float]] = mapped_column(HALFVEC(1024))

from pgvector.sqlalchemy import SPARSEVEC

class Document(Base):
    sparse: Mapped[SparseVec] = mapped_column(SPARSEVEC(10000))

@event.listens_for(Document, "before_update")
def update_embedding(mapper, conn, target):
    if target.body_changed:
        target.embedding = compute_embedding_sync(target.body)

from sqlalchemy.dialects.postgresql import insert

# Two separate queries; combine with RRF in app
fts_results = await session.execute(
    select(Document.id).where(Document.search_vec.op("@@")(q_ts)).limit(50)
)
vec_results = await session.execute(
    select(Document.id).order_by(Document.embedding.cosine_distance(q_emb)).limit(50)
)

def rrf(rankings, k=60):
    scores = defaultdict(float)
    for ranking in rankings:
        for rank, item_id in enumerate(ranking):
            scores[item_id] += 1 / (k + rank)
    return sorted(scores.items(), key=lambda x: -x[1])