pgvector

Store and query vectors in Postgres using the pgvector extension.

pgvector is a Postgres extension that adds a vector column type with cosine, L2, and inner-product operators. It rides on whatever Postgres you already have — including the one Flux uses for its own event log — so a small RAG workflow can run with zero new infrastructure.

Install the extension

pgvector is server-side. Install it on your Postgres instance, then enable it once per database:

CREATE EXTENSION IF NOT EXISTS vector;

Most managed Postgres providers (RDS, Cloud SQL, Supabase, Neon) include pgvector in their default extension list — CREATE EXTENSION is the only step on those. For self-hosted Postgres you may need to install the postgresql-NN-pgvector package first.

Driver

Flux 0.56.0 uses psycopg (v3) for its own Postgres connection (pyproject.toml, postgresql extra pulls psycopg[binary,pool]). The same driver covers both sync and async task code, so one install serves everything:

pip install "psycopg[binary,pool]" pgvector

If the worker already has the postgresql extra installed (pip install 'flux-core[postgresql]'), only pgvector is missing.

Schema

CREATE TABLE documents (
    id          bigserial PRIMARY KEY,
    text        text NOT NULL,
    embedding   vector(1536),
    metadata    jsonb,
    created_at  timestamptz DEFAULT now()
);

Pick vector(N) to match your embedding model — 1536 for OpenAI text-embedding-3-small, 1024 for Cohere v3, 384 for all-MiniLM-L6-v2. The dimension is part of the column type; changing it requires ALTER TABLE.

Index

pgvector ships two index types. HNSW is the right default for cosine similarity at scale:

CREATE INDEX ON documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

HNSW builds slower than IVFFlat but gives better recall at low query latency and doesn’t need a per-table training step. Use IVFFlat (vector_l2_ops, vector_cosine_ops, vector_ip_ops) only when you have hundreds of millions of rows and need the smaller index footprint.

Tasks

import json
from flux import workflow, task, ExecutionContext


@task.with_options(secret_requests=["pgvector_dsn"])
async def add_documents(secrets, docs: list[dict]) -> int:
    import psycopg
    from pgvector.psycopg import register_vector

    async with await psycopg.AsyncConnection.connect(secrets["pgvector_dsn"]) as conn:
        await register_vector(conn)
        async with conn.cursor() as cur:
            await cur.executemany(
                "INSERT INTO documents (text, embedding, metadata) VALUES (%s, %s, %s)",
                [(d["text"], d["embedding"], json.dumps(d.get("metadata", {}))) for d in docs],
            )
        await conn.commit()
    return len(docs)


@task.with_options(secret_requests=["pgvector_dsn"])
async def search(secrets, embedding: list[float], k: int = 5) -> list[dict]:
    import psycopg
    from pgvector.psycopg import register_vector

    async with await psycopg.AsyncConnection.connect(secrets["pgvector_dsn"]) as conn:
        await register_vector(conn)
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT id, text, metadata, embedding <=> %s::vector AS distance "
                "FROM documents ORDER BY embedding <=> %s::vector LIMIT %s",
                (embedding, embedding, k),
            )
            rows = await cur.fetchall()
    return [
        {"id": r[0], "text": r[1], "metadata": r[2], "distance": float(r[3])}
        for r in rows
    ]

<=> is cosine distance, <-> is L2, <#> is negative inner product. Match the operator to the *_ops you used when creating the index, otherwise Postgres falls back to a sequential scan.

Tempting and usually fine: put your documents table in the same database as Flux’s event log, in a different schema (CREATE SCHEMA rag; SET search_path TO rag, public;). Pros: one less thing to back up, one connection pool. Cons: a long vector query holds a connection, and it competes with Flux’s own pool (pool_size=20, max_overflow=20 per process by default). If you embed-and-query at high throughput, give the vector workload its own database to avoid starving the workflow engine.

What goes wrong

Forgetting register_vector(conn). Without it, psycopg sees vector columns as plain text and returns string representations of arrays.
Wrong index operator class. vector_l2_ops index + <=> query = no index usage. EXPLAIN ANALYZE will say so.
Dimension drift. Same trap as every vector store: re-embedding with a new model invalidates the whole table. Either version the table (documents_v2) or include the model name in metadata and filter at query time.