pgvector-python

pgvector support for Python

Supports Django, SQLAlchemy, SQLModel, Psycopg 3, Psycopg 2, asyncpg, and Peewee

Installation

Run:

pip install pgvector

And follow the instructions for your database library:

Django
SQLAlchemy
SQLModel
Psycopg 3
Psycopg 2
asyncpg
Peewee

Or check out some examples:

Embeddings with OpenAI
Binary embeddings with Cohere
Sentence embeddings with SentenceTransformers
Hybrid search with SentenceTransformers (Reciprocal Rank Fusion)
Hybrid search with SentenceTransformers (cross-encoder)
Sparse search with Transformers
Late interaction search with ColBERT
Visual document retrieval with ColPali
Image search with PyTorch
Image search with perceptual hashing
Morgan fingerprints with RDKit
Topic modeling with Gensim
Implicit feedback recommendations with Implicit
Explicit feedback recommendations with Surprise
Recommendations with LightFM
Horizontal scaling with Citus
Bulk loading with COPY

Django

Create a migration to enable the extension

from pgvector.django import VectorExtension

class Migration(migrations.Migration):
    operations = [
        VectorExtension()
    ]

Add a vector field to your model

from pgvector.django import VectorField

class Item(models.Model):
    embedding = VectorField(dimensions=3)

Also supports HalfVectorField, BitField, and SparseVectorField

Insert a vector

item = Item(embedding=[1, 2, 3])
item.save()

Get the nearest neighbors to a vector

from pgvector.django import L2Distance

Item.objects.order_by(L2Distance('embedding', [3, 1, 2]))[:5]

Also supports MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, and JaccardDistance

Get the distance

Item.objects.annotate(distance=L2Distance('embedding', [3, 1, 2]))

Get items within a certain distance

Item.objects.alias(distance=L2Distance('embedding', [3, 1, 2])).filter(distance__lt=5)

Average vectors

from django.db.models import Avg

Item.objects.aggregate(Avg('embedding'))

Also supports Sum

Add an approximate index

from pgvector.django import HnswIndex, IvfflatIndex

class Item(models.Model):
    class Meta:
        indexes = [
            HnswIndex(
                name='my_index',
                fields=['embedding'],
                m=16,
                ef_construction=64,
                opclasses=['vector_l2_ops']
            ),
            # or
            IvfflatIndex(
                name='my_index',
                fields=['embedding'],
                lists=100,
                opclasses=['vector_l2_ops']
            )
        ]

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

Half-Precision Indexing

Index vectors at half-precision

from django.contrib.postgres.indexes import OpClass
from django.db.models.functions import Cast
from pgvector.django import HnswIndex, HalfVectorField

class Item(models.Model):
    class Meta:
        indexes = [
            HnswIndex(
                OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'),
                name='my_index',
                m=16,
                ef_construction=64
            )
        ]

Note: Add 'django.contrib.postgres' to INSTALLED_APPS to use OpClass

Get the nearest neighbors

distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [3, 1, 2])
Item.objects.order_by(distance)[:5]

SQLAlchemy

Enable the extension

session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))

Add a vector column

from pgvector.sqlalchemy import Vector

class Item(Base):
    embedding = mapped_column(Vector(3))

Also supports HALFVEC, BIT, and SPARSEVEC

Insert a vector

item = Item(embedding=[1, 2, 3])
session.add(item)
session.commit()

Get the nearest neighbors to a vector

session.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5))

Also supports max_inner_product, cosine_distance, l1_distance, hamming_distance, and jaccard_distance

Get the distance

session.scalars(select(Item.embedding.l2_distance([3, 1, 2])))

Get items within a certain distance

session.scalars(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5))

Average vectors

from pgvector.sqlalchemy import avg

session.scalars(select(avg(Item.embedding))).first()

Also supports sum

Add an approximate index

index = Index(
    'my_index',
    Item.embedding,
    postgresql_using='hnsw',
    postgresql_with={'m': 16, 'ef_construction': 64},
    postgresql_ops={'embedding': 'vector_l2_ops'}
)
# or
index = Index(
    'my_index',
    Item.embedding,
    postgresql_using='ivfflat',
    postgresql_with={'lists': 100},
    postgresql_ops={'embedding': 'vector_l2_ops'}
)

index.create(engine)

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

Half-Precision Indexing

Index vectors at half-precision

from pgvector.sqlalchemy import HALFVEC
from sqlalchemy.sql import func

index = Index(
    'my_index',
    func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
    postgresql_using='hnsw',
    postgresql_with={'m': 16, 'ef_construction': 64},
    postgresql_ops={'embedding': 'halfvec_l2_ops'}
)

Get the nearest neighbors

order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2])
session.scalars(select(Item).order_by(order).limit(5))

Arrays

Add an array column

from pgvector.sqlalchemy import Vector
from sqlalchemy import ARRAY

class Item(Base):
    embeddings = mapped_column(ARRAY(Vector(3)))

And register the types with the underlying driver

For Psycopg 3, use

from pgvector.psycopg import register_vector
from sqlalchemy import event

@event.listens_for(engine, "connect")
def connect(dbapi_connection, connection_record):
    register_vector(dbapi_connection)

For async connections with Psycopg 3, use

from pgvector.psycopg import register_vector_async
from sqlalchemy import event

@event.listens_for(engine.sync_engine, "connect")
def connect(dbapi_connection, connection_record):
    dbapi_connection.run_async(register_vector_async)

For Psycopg 2, use

from pgvector.psycopg2 import register_vector
from sqlalchemy import event

@event.listens_for(engine, "connect")
def connect(dbapi_connection, connection_record):
    register_vector(dbapi_connection, arrays=True)

SQLModel

Enable the extension

session.exec(text('CREATE EXTENSION IF NOT EXISTS vector'))

Add a vector column

from pgvector.sqlalchemy import Vector
from sqlalchemy import Column

class Item(SQLModel, table=True):
    embedding: Any = Field(sa_column=Column(Vector(3)))

Also supports HALFVEC, BIT, and SPARSEVEC

Insert a vector

item = Item(embedding=[1, 2, 3])
session.add(item)
session.commit()

Get the nearest neighbors to a vector

session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5))

Also supports max_inner_product, cosine_distance, l1_distance, hamming_distance, and jaccard_distance

Get the distance

session.exec(select(Item.embedding.l2_distance([3, 1, 2])))

Get items within a certain distance

session.exec(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5))

Average vectors

from pgvector.sqlalchemy import avg

session.exec(select(avg(Item.embedding))).first()

Also supports sum

Add an approximate index

from sqlalchemy import Index

index = Index(
    'my_index',
    Item.embedding,
    postgresql_using='hnsw',
    postgresql_with={'m': 16, 'ef_construction': 64},
    postgresql_ops={'embedding': 'vector_l2_ops'}
)
# or
index = Index(
    'my_index',
    Item.embedding,
    postgresql_using='ivfflat',
    postgresql_with={'lists': 100},
    postgresql_ops={'embedding': 'vector_l2_ops'}
)

index.create(engine)

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

Psycopg 3

Enable the extension

conn.execute('CREATE EXTENSION IF NOT EXISTS vector')

Register the vector type with your connection

from pgvector.psycopg import register_vector

register_vector(conn)

For connection pools, use

def configure(conn):
    register_vector(conn)

pool = ConnectionPool(..., configure=configure)

For async connections, use

from pgvector.psycopg import register_vector_async

await register_vector_async(conn)

Create a table

conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')

Insert a vector

embedding = np.array([1, 2, 3])
conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))

Get the nearest neighbors to a vector

conn.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall()

Add an approximate index

conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)')
# or
conn.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)')

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

Psycopg 2

Enable the extension

cur = conn.cursor()
cur.execute('CREATE EXTENSION IF NOT EXISTS vector')

Register the vector type with your connection or cursor

from pgvector.psycopg2 import register_vector

register_vector(conn)

Create a table

cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')

Insert a vector

embedding = np.array([1, 2, 3])
cur.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))

Get the nearest neighbors to a vector

cur.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,))
cur.fetchall()

Add an approximate index

cur.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)')
# or
cur.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)')

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

asyncpg

Enable the extension

await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')

Register the vector type with your connection

from pgvector.asyncpg import register_vector

await register_vector(conn)

or your pool

async def init(conn):
    await register_vector(conn)

pool = await asyncpg.create_pool(..., init=init)

Create a table

await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')

Insert a vector

embedding = np.array([1, 2, 3])
await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding)

Get the nearest neighbors to a vector

await conn.fetch('SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5', embedding)

Add an approximate index

await conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)')
# or
await conn.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)')

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

Peewee

Add a vector column

from pgvector.peewee import VectorField

class Item(BaseModel):
    embedding = VectorField(dimensions=3)

Also supports HalfVectorField, FixedBitField, and SparseVectorField

Insert a vector

item = Item.create(embedding=[1, 2, 3])

Get the nearest neighbors to a vector

Item.select().order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)

Also supports max_inner_product, cosine_distance, l1_distance, hamming_distance, and jaccard_distance

Get the distance

Item.select(Item.embedding.l2_distance([3, 1, 2]).alias('distance'))

Get items within a certain distance

Item.select().where(Item.embedding.l2_distance([3, 1, 2]) < 5)

Average vectors

from peewee import fn

Item.select(fn.avg(Item.embedding).coerce(True)).scalar()

Also supports sum

Add an approximate index

Item.add_index('embedding vector_l2_ops', using='hnsw')

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

History

View the changelog

Contributing

Everyone is encouraged to help improve this project. Here are a few ways you can help:

Report bugs
Fix bugs and submit pull requests
Write, clarify, or fix documentation
Suggest or add new features

To get started with development:

git clone https://github.com/pgvector/pgvector-python.git
cd pgvector-python
pip install -r requirements.txt
createdb pgvector_python_test
pytest

To run an example:

cd examples/loading
pip install -r requirements.txt
createdb pgvector_example
python3 example.py

License

razi-tm/pgvector-python

ErrorLooks like something went wrong!

About

Resources

License

Uh oh!

Stars

Watchers

Forks

Releases

Packages 0

Languages

Error
Looks like something went wrong!

Packages