diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 4d4e8ed..85db5b6 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,20 +6,20 @@ jobs: strategy: fail-fast: false matrix: - python: [3.13, 3.9] + python: [3.14, 3.9] steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v5 + - uses: actions/setup-python@v6 with: python-version: ${{ matrix.python }} - - run: pip install -r requirements.txt + - run: pip install --group dev - uses: ankane/setup-postgres@v1 with: database: pgvector_python_test dev-files: true - run: | cd /tmp - git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git + git clone --branch v0.8.2 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install diff --git a/.gitignore b/.gitignore index f7ff659..5556c9f 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,5 @@ venv/ *.pyc __pycache__ .pytest_cache/ +*.lock +examples/rag/README.md diff --git a/CHANGELOG.md b/CHANGELOG.md index ebc165a..745335f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,17 @@ -## 0.4.0 (unreleased) +## 0.4.2 (2025-12-04) + +- Added support for Django 6 +- Added support for `str` objects for `bit` type with SQLAlchemy + +## 0.4.1 (2025-04-26) + +- Fixed `SparseVector` constructor for SciPy sparse matrices + +## 0.4.0 (2025-03-15) - Added top-level `pgvector` package - Added support for pg8000 +- Added support for `bytes` to `Bit` constructor - Changed `globally` option to default to `False` for Psycopg 2 - Changed `arrays` option to default to `True` for Psycopg 2 - Fixed equality for `Vector`, `HalfVector`, `Bit`, and `SparseVector` classes diff --git a/README.md b/README.md index 299753e..95d5fbe 100644 --- a/README.md +++ b/README.md @@ -22,11 +22,12 @@ And follow the instructions for your database library: - [Psycopg 3](#psycopg-3) - [Psycopg 2](#psycopg-2) - [asyncpg](#asyncpg) -- [pg8000](#pg8000) [unreleased] +- [pg8000](#pg8000) - [Peewee](#peewee) Or check out some examples: +- [Retrieval-augmented generation](https://github.com/pgvector/pgvector-python/blob/master/examples/rag/example.py) with Ollama - [Embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/openai/example.py) with OpenAI - [Binary embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/cohere/example.py) with Cohere - [Sentence embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/sentence_transformers/example.py) with SentenceTransformers @@ -176,10 +177,10 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) Add a vector column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR class Item(Base): - embedding = mapped_column(Vector(3)) + embedding = mapped_column(VECTOR(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` @@ -258,7 +259,6 @@ 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'} ) ``` @@ -270,16 +270,47 @@ order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2]) session.scalars(select(Item).order_by(order).limit(5)) ``` +#### Binary Quantization + +Use expression indexing for binary quantization + +```python +from pgvector.sqlalchemy import BIT +from sqlalchemy.sql import func + +index = Index( + 'my_index', + func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'), + postgresql_using='hnsw', + postgresql_ops={'embedding': 'bit_hamming_ops'} +) +``` + +Get the nearest neighbors by Hamming distance + +```python +order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) +session.scalars(select(Item).order_by(order).limit(5)) +``` + +Re-rank by the original vectors for better recall + +```python +order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) +subquery = session.query(Item).order_by(order).limit(20).subquery() +session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5)) +``` + #### Arrays Add an array column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import ARRAY class Item(Base): - embeddings = mapped_column(ARRAY(Vector(3))) + embeddings = mapped_column(ARRAY(VECTOR(3))) ``` And register the types with the underlying driver @@ -314,7 +345,7 @@ from sqlalchemy import event @event.listens_for(engine, "connect") def connect(dbapi_connection, connection_record): - register_vector(dbapi_connection, arrays=True) + register_vector(dbapi_connection) ``` ## SQLModel @@ -328,10 +359,10 @@ session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) Add a vector column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR class Item(SQLModel, table=True): - embedding: Any = Field(sa_type=Vector(3)) + embedding: Any = Field(sa_type=VECTOR(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` @@ -408,7 +439,7 @@ Enable the extension conn.execute('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection +Register the types with your connection ```python from pgvector.psycopg import register_vector @@ -471,7 +502,7 @@ cur = conn.cursor() cur.execute('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection or cursor +Register the types with your connection or cursor ```python from pgvector.psycopg2 import register_vector @@ -517,7 +548,7 @@ Enable the extension await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection +Register the types with your connection ```python from pgvector.asyncpg import register_vector @@ -571,7 +602,7 @@ Enable the extension conn.run('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection +Register the types with your connection ```python from pgvector.pg8000 import register_vector @@ -776,7 +807,7 @@ To get started with development: ```sh git clone https://github.com/pgvector/pgvector-python.git cd pgvector-python -pip install -r requirements.txt +pip install --group dev createdb pgvector_python_test pytest ``` @@ -785,7 +816,7 @@ To run an example: ```sh cd examples/loading -pip install -r requirements.txt +pip install --group dev createdb pgvector_example python3 example.py ``` diff --git a/examples/citus/pyproject.toml b/examples/citus/pyproject.toml new file mode 100644 index 0000000..ee40a36 --- /dev/null +++ b/examples/citus/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/citus/requirements.txt b/examples/citus/requirements.txt deleted file mode 100644 index 1cf8ee9..0000000 --- a/examples/citus/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -numpy -pgvector -psycopg[binary] diff --git a/examples/cohere/example.py b/examples/cohere/example.py index 393d1e0..5ef4eec 100644 --- a/examples/cohere/example.py +++ b/examples/cohere/example.py @@ -9,12 +9,12 @@ register_vector(conn) conn.execute('DROP TABLE IF EXISTS documents') -conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1536))') def embed(input, input_type): - co = cohere.Client() - response = co.embed(texts=input, model='embed-english-v3.0', input_type=input_type, embedding_types=['ubinary']) + co = cohere.ClientV2() + response = co.embed(texts=input, model='embed-v4.0', input_type=input_type, embedding_types=['ubinary']) return [np.unpackbits(np.array(embedding, dtype=np.uint8)) for embedding in response.embeddings.ubinary] diff --git a/examples/cohere/pyproject.toml b/examples/cohere/pyproject.toml new file mode 100644 index 0000000..f0c88b7 --- /dev/null +++ b/examples/cohere/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "cohere", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/cohere/requirements.txt b/examples/cohere/requirements.txt deleted file mode 100644 index 22fd056..0000000 --- a/examples/cohere/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -cohere -pgvector -psycopg[binary] diff --git a/examples/colbert/approximate.py b/examples/colbert/approximate.py new file mode 100644 index 0000000..41f88b2 --- /dev/null +++ b/examples/colbert/approximate.py @@ -0,0 +1,80 @@ +# based on section 3.6 of https://arxiv.org/abs/2004.12832 + +from colbert.infra import ColBERTConfig +from colbert.modeling.checkpoint import Checkpoint +from pgvector.psycopg import register_vector +import psycopg +import warnings + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('DROP TABLE IF EXISTS document_embeddings') + +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text)') +conn.execute('CREATE TABLE document_embeddings (id bigserial PRIMARY KEY, document_id bigint, embedding vector(128))') + +conn.execute(""" +CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$ + WITH queries AS ( + SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) + ), + documents AS ( + SELECT unnest(document) AS document + ), + similarities AS ( + SELECT query_number, 1 - (document <=> query) AS similarity FROM queries CROSS JOIN documents + ), + max_similarities AS ( + SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number + ) + SELECT SUM(max_similarity) FROM max_similarities +$$ LANGUAGE SQL +""") + +warnings.filterwarnings('ignore') # ignore warnings from colbert + +config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) +checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +doc_embeddings = checkpoint.docFromText(input, keep_dims=False) +for content, embeddings in zip(input, doc_embeddings): + with conn.transaction(): + result = conn.execute('INSERT INTO documents (content) VALUES (%s) RETURNING id', (content,)).fetchone() + params = [] + for embedding in embeddings: + params.extend([result[0], embedding.numpy()]) + values = ', '.join(['(%s, %s)' for _ in embeddings]) + conn.execute(f'INSERT INTO document_embeddings (document_id, embedding) VALUES {values}', params) + +conn.execute('CREATE INDEX ON document_embeddings (document_id)') +conn.execute('CREATE INDEX ON document_embeddings USING hnsw (embedding vector_cosine_ops)') + +query = 'puppy' +query_embeddings = [e.numpy() for e in checkpoint.queryFromText([query])[0]] +approximate_stage = ' UNION ALL '.join(['(SELECT document_id FROM document_embeddings ORDER BY embedding <=> %s LIMIT 5)' for _ in query_embeddings]) +sql = f""" +WITH approximate_stage AS ( + {approximate_stage} +), +embeddings AS ( + SELECT document_id, array_agg(embedding) AS embeddings FROM document_embeddings + WHERE document_id IN (SELECT DISTINCT document_id FROM approximate_stage) + GROUP BY document_id +) +SELECT content, max_sim(embeddings, %s) AS max_sim FROM documents +INNER JOIN embeddings ON embeddings.document_id = documents.id +ORDER BY max_sim DESC LIMIT 10 +""" +params = query_embeddings + [query_embeddings] +result = conn.execute(sql, params).fetchall() +for row in result: + print(row) diff --git a/examples/colbert/exact.py b/examples/colbert/exact.py index 1c90b47..e6a2936 100644 --- a/examples/colbert/exact.py +++ b/examples/colbert/exact.py @@ -2,6 +2,7 @@ from colbert.modeling.checkpoint import Checkpoint from pgvector.psycopg import register_vector import psycopg +import warnings conn = psycopg.connect(dbname='pgvector_example', autocommit=True) @@ -28,6 +29,8 @@ $$ LANGUAGE SQL """) +warnings.filterwarnings('ignore') # ignore warnings from colbert + config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) diff --git a/examples/colbert/pyproject.toml b/examples/colbert/pyproject.toml new file mode 100644 index 0000000..face4d2 --- /dev/null +++ b/examples/colbert/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "colbert-ai", + "pgvector", + "psycopg[binary]", + "transformers==4.49.0" +] diff --git a/examples/colbert/requirements.txt b/examples/colbert/requirements.txt deleted file mode 100644 index 4402ce8..0000000 --- a/examples/colbert/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -colbert-ai -pgvector -psycopg[binary] diff --git a/examples/colpali/pyproject.toml b/examples/colpali/pyproject.toml new file mode 100644 index 0000000..23fb23f --- /dev/null +++ b/examples/colpali/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "colpali-engine", + "datasets", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/colpali/requirements.txt b/examples/colpali/requirements.txt deleted file mode 100644 index 4cf770d..0000000 --- a/examples/colpali/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -colpali-engine -datasets -pgvector -psycopg[binary] diff --git a/examples/gensim/pyproject.toml b/examples/gensim/pyproject.toml new file mode 100644 index 0000000..7a33423 --- /dev/null +++ b/examples/gensim/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "gensim", + "numpy", + "pgvector", + "psycopg[binary]", + "scipy<1.13" +] diff --git a/examples/gensim/requirements.txt b/examples/gensim/requirements.txt deleted file mode 100644 index 15411cd..0000000 --- a/examples/gensim/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -gensim -numpy -pgvector -psycopg[binary] -scipy<1.13 diff --git a/examples/hybrid_search/pyproject.toml b/examples/hybrid_search/pyproject.toml new file mode 100644 index 0000000..b5a904a --- /dev/null +++ b/examples/hybrid_search/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "sentence-transformers" +] diff --git a/examples/hybrid_search/requirements.txt b/examples/hybrid_search/requirements.txt deleted file mode 100644 index 237dcd1..0000000 --- a/examples/hybrid_search/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -pgvector -psycopg[binary] -sentence-transformers diff --git a/examples/image_search/pyproject.toml b/examples/image_search/pyproject.toml new file mode 100644 index 0000000..7644382 --- /dev/null +++ b/examples/image_search/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "matplotlib", + "pgvector", + "psycopg[binary]", + "torch", + "torchvision", + "tqdm" +] diff --git a/examples/image_search/requirements.txt b/examples/image_search/requirements.txt deleted file mode 100644 index 3d82365..0000000 --- a/examples/image_search/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ -matplotlib -pgvector -psycopg[binary] -torch -torchvision -tqdm diff --git a/examples/imagehash/pyproject.toml b/examples/imagehash/pyproject.toml new file mode 100644 index 0000000..cf06c2b --- /dev/null +++ b/examples/imagehash/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "datasets", + "imagehash", + "matplotlib", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/imagehash/requirements.txt b/examples/imagehash/requirements.txt deleted file mode 100644 index e3971e6..0000000 --- a/examples/imagehash/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -datasets -imagehash -matplotlib -pgvector -psycopg[binary] diff --git a/examples/implicit/example.py b/examples/implicit/example.py index f70eb8c..2cbf7c6 100644 --- a/examples/implicit/example.py +++ b/examples/implicit/example.py @@ -1,6 +1,6 @@ import implicit from implicit.datasets.movielens import get_movielens -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Integer, String from sqlalchemy.orm import declarative_base, mapped_column, Session @@ -16,7 +16,7 @@ class User(Base): __tablename__ = 'user' id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) class Item(Base): @@ -24,7 +24,7 @@ class Item(Base): id = mapped_column(Integer, primary_key=True) title = mapped_column(String) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) Base.metadata.drop_all(engine) diff --git a/examples/implicit/pyproject.toml b/examples/implicit/pyproject.toml new file mode 100644 index 0000000..c03b187 --- /dev/null +++ b/examples/implicit/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "h5py", + "implicit", + "pgvector", + "psycopg[binary]", + "SQLAlchemy" +] diff --git a/examples/implicit/requirements.txt b/examples/implicit/requirements.txt deleted file mode 100644 index 424abbd..0000000 --- a/examples/implicit/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -h5py -implicit -pgvector -psycopg[binary] -SQLAlchemy diff --git a/examples/lightfm/example.py b/examples/lightfm/example.py index fcb9027..65031c4 100644 --- a/examples/lightfm/example.py +++ b/examples/lightfm/example.py @@ -1,6 +1,6 @@ from lightfm import LightFM from lightfm.datasets import fetch_movielens -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Float, Integer, String from sqlalchemy.orm import declarative_base, mapped_column, Session @@ -16,7 +16,7 @@ class User(Base): __tablename__ = 'user' id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) class Item(Base): @@ -24,7 +24,7 @@ class Item(Base): id = mapped_column(Integer, primary_key=True) title = mapped_column(String) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) bias = mapped_column(Float) diff --git a/examples/lightfm/pyproject.toml b/examples/lightfm/pyproject.toml new file mode 100644 index 0000000..c202058 --- /dev/null +++ b/examples/lightfm/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "lightfm", + "pgvector", + "psycopg[binary]", + "SQLAlchemy" +] diff --git a/examples/lightfm/requirements.txt b/examples/lightfm/requirements.txt deleted file mode 100644 index cfa5f51..0000000 --- a/examples/lightfm/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -lightfm -pgvector -psycopg[binary] -SQLAlchemy diff --git a/examples/loading/example.py b/examples/loading/example.py index 0702129..7f3dce8 100644 --- a/examples/loading/example.py +++ b/examples/loading/example.py @@ -25,12 +25,12 @@ copy.set_types(['vector']) for i, embedding in enumerate(embeddings): + copy.write_row([embedding]) + # show progress if i % 10000 == 0: print('.', end='', flush=True) - copy.write_row([embedding]) - print('\nSuccess!') # create any indexes *after* loading initial data (skipping for this example) diff --git a/examples/loading/pyproject.toml b/examples/loading/pyproject.toml new file mode 100644 index 0000000..ee40a36 --- /dev/null +++ b/examples/loading/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/loading/requirements.txt b/examples/loading/requirements.txt deleted file mode 100644 index 1cf8ee9..0000000 --- a/examples/loading/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -numpy -pgvector -psycopg[binary] diff --git a/examples/openai/pyproject.toml b/examples/openai/pyproject.toml new file mode 100644 index 0000000..3e6661a --- /dev/null +++ b/examples/openai/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "openai", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/openai/requirements.txt b/examples/openai/requirements.txt deleted file mode 100644 index 18587e2..0000000 --- a/examples/openai/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -openai -pgvector -psycopg[binary] diff --git a/examples/rag/example.py b/examples/rag/example.py new file mode 100644 index 0000000..4d5d307 --- /dev/null +++ b/examples/rag/example.py @@ -0,0 +1,65 @@ +# Run: +# ollama pull llama3.2 +# ollama pull nomic-embed-text +# ollama serve + +import numpy as np +import ollama +from pathlib import Path +from pgvector.psycopg import register_vector +import psycopg +import urllib.request + +query = 'What index types are supported?' +load_data = True + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +if load_data: + # get data + url = 'https://raw.githubusercontent.com/pgvector/pgvector/refs/heads/master/README.md' + dest = Path(__file__).parent / 'README.md' + if not dest.exists(): + urllib.request.urlretrieve(url, dest) + + with open(dest, encoding='utf-8') as f: + doc = f.read() + + # generate chunks + # TODO improve chunking + # TODO remove markdown + chunks = doc.split('\n## ') + + # embed chunks + # nomic-embed-text has task instruction prefix + input = ['search_document: ' + chunk for chunk in chunks] + embeddings = ollama.embed(model='nomic-embed-text', input=input).embeddings + + # create table + conn.execute('DROP TABLE IF EXISTS chunks') + conn.execute('CREATE TABLE chunks (id bigserial PRIMARY KEY, content text, embedding vector(768))') + + # store chunks + cur = conn.cursor() + with cur.copy('COPY chunks (content, embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: + copy.set_types(['text', 'vector']) + + for content, embedding in zip(chunks, embeddings): + copy.write_row([content, embedding]) + +# embed query +# nomic-embed-text has task instruction prefix +input = 'search_query: ' + query +embedding = ollama.embed(model='nomic-embed-text', input=input).embeddings[0] + +# retrieve chunks +result = conn.execute('SELECT content FROM chunks ORDER BY embedding <=> %s LIMIT 5', (np.array(embedding),)).fetchall() +context = '\n\n'.join([row[0] for row in result]) + +# get answer +# TODO improve prompt +prompt = f'Answer this question: {query}\n\n{context}' +response = ollama.generate(model='llama3.2', prompt=prompt).response +print(response) diff --git a/examples/rag/pyproject.toml b/examples/rag/pyproject.toml new file mode 100644 index 0000000..fa0dcfd --- /dev/null +++ b/examples/rag/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "ollama", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/rdkit/pyproject.toml b/examples/rdkit/pyproject.toml new file mode 100644 index 0000000..f8c035a --- /dev/null +++ b/examples/rdkit/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "rdkit" +] diff --git a/examples/rdkit/requirements.txt b/examples/rdkit/requirements.txt deleted file mode 100644 index 85a3e4f..0000000 --- a/examples/rdkit/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -pgvector -psycopg[binary] -rdkit diff --git a/examples/sentence_transformers/pyproject.toml b/examples/sentence_transformers/pyproject.toml new file mode 100644 index 0000000..b5a904a --- /dev/null +++ b/examples/sentence_transformers/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "sentence-transformers" +] diff --git a/examples/sentence_transformers/requirements.txt b/examples/sentence_transformers/requirements.txt deleted file mode 100644 index 237dcd1..0000000 --- a/examples/sentence_transformers/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -pgvector -psycopg[binary] -sentence-transformers diff --git a/examples/sparse_search/pyproject.toml b/examples/sparse_search/pyproject.toml new file mode 100644 index 0000000..7927c34 --- /dev/null +++ b/examples/sparse_search/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]", + "torch", + "transformers" +] diff --git a/examples/sparse_search/requirements.txt b/examples/sparse_search/requirements.txt deleted file mode 100644 index 3de81c7..0000000 --- a/examples/sparse_search/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -numpy -pgvector -psycopg[binary] -torch -transformers diff --git a/examples/surprise/example.py b/examples/surprise/example.py index bd7d18d..e413bcf 100644 --- a/examples/surprise/example.py +++ b/examples/surprise/example.py @@ -1,4 +1,4 @@ -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Integer from sqlalchemy.orm import declarative_base, mapped_column, Session from surprise import Dataset, SVD @@ -15,14 +15,14 @@ class User(Base): __tablename__ = 'user' id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) class Item(Base): __tablename__ = 'item' id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) Base.metadata.drop_all(engine) diff --git a/examples/surprise/pyproject.toml b/examples/surprise/pyproject.toml new file mode 100644 index 0000000..94c6f13 --- /dev/null +++ b/examples/surprise/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "scikit-surprise", + "SQLAlchemy" +] diff --git a/examples/surprise/requirements.txt b/examples/surprise/requirements.txt deleted file mode 100644 index cb2dca4..0000000 --- a/examples/surprise/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -pgvector -psycopg[binary] -scikit-surprise -SQLAlchemy diff --git a/pgvector/bit.py b/pgvector/bit.py index 4be7385..26a9d8d 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -1,54 +1,64 @@ import numpy as np from struct import pack, unpack_from +from warnings import warn class Bit: def __init__(self, value): - if isinstance(value, str): - self._value = self.from_text(value)._value + if isinstance(value, bytes): + self._len = 8 * len(value) + self._data = value else: - if isinstance(value, np.ndarray): - if value.dtype == np.uint8: - value = np.unpackbits(value).astype(bool) - elif value.dtype != np.bool: - raise ValueError('expected dtype to be bool or uint8') + if isinstance(value, str): + value = [v != '0' for v in value] else: - value = np.asarray(value, dtype=bool) + value = np.asarray(value) - if value.ndim != 1: - raise ValueError('expected ndim to be 1') + if value.dtype != np.bool: + # skip warning for result of np.unpackbits + if value.dtype != np.uint8 or np.any(value > 1): + warn('expected elements to be boolean', stacklevel=2) + value = value.astype(bool) - self._value = value + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + self._len = len(value) + self._data = np.packbits(value).tobytes() def __repr__(self): return f'Bit({self.to_text()})' def __eq__(self, other): if isinstance(other, self.__class__): - return np.array_equal(self.to_numpy(), other.to_numpy()) + return self._len == other._len and self._data == other._data return False def to_list(self): - return self._value.tolist() + return self.to_numpy().tolist() def to_numpy(self): - return self._value + return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool) def to_text(self): - return ''.join(self._value.astype(np.uint8).astype(str)) + return ''.join(format(v, '08b') for v in self._data)[:self._len] def to_binary(self): - return pack('>i', len(self._value)) + np.packbits(self._value).tobytes() + return pack('>i', self._len) + self._data @classmethod def from_text(cls, value): - return cls(np.asarray([v != '0' for v in value], dtype=bool)) + return cls(str(value)) @classmethod def from_binary(cls, value): - count = unpack_from('>i', value)[0] - buf = np.frombuffer(value, dtype=np.uint8, offset=4) - return cls(np.unpackbits(buf, count=count).astype(bool)) + if not isinstance(value, bytes): + raise ValueError('expected bytes') + + bit = cls.__new__(cls) + bit._len = unpack_from('>i', value)[0] + bit._data = value[4:] + return bit @classmethod def _to_db(cls, value): diff --git a/pgvector/django/extensions.py b/pgvector/django/extensions.py index 0573f72..1d04739 100644 --- a/pgvector/django/extensions.py +++ b/pgvector/django/extensions.py @@ -1,6 +1,11 @@ +from django import VERSION from django.contrib.postgres.operations import CreateExtension class VectorExtension(CreateExtension): - def __init__(self): - self.name = 'vector' + if VERSION[0] >= 6: + def __init__(self, hints=None): + super().__init__('vector', hints=hints) + else: + def __init__(self): + self.name = 'vector' diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index b40c673..33e5124 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,11 +1,10 @@ from .register import register_vector # TODO remove -from .. import HalfVector, SparseVector, Vector +from .. import HalfVector, SparseVector __all__ = [ 'register_vector', - 'Vector', 'HalfVector', 'SparseVector' ] diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 8df2dfd..895fbd0 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -85,7 +85,7 @@ def _from_sparse(self, value): if hasattr(value, 'coords'): # scipy 1.13+ - self._indices = value.coords[0].tolist() + self._indices = value.coords[-1].tolist() else: self._indices = value.col.tolist() self._values = value.data.tolist() diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 0f83f3c..1ea85c3 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -14,6 +14,18 @@ def get_col_spec(self, **kw): return 'BIT' return 'BIT(%d)' % self.length + def bind_processor(self, dialect): + if dialect.__class__.__name__ == 'PGDialect_asyncpg': + import asyncpg + + def process(value): + if isinstance(value, str): + return asyncpg.BitString(value) + return value + return process + else: + return super().bind_processor(dialect) + class comparator_factory(UserDefinedType.Comparator): def hamming_distance(self, other): return self.op('<~>', return_type=Float)(other) diff --git a/pyproject.toml b/pyproject.toml index 0f291f5..5716d05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta" [project] name = "pgvector" -version = "0.3.6" +version = "0.4.2" description = "pgvector support for Python" readme = "README.md" authors = [ {name = "Andrew Kane", email = "andrew@ankane.org"} ] -license = {text = "MIT"} +license = "MIT" requires-python = ">= 3.9" dependencies = [ "numpy" @@ -19,6 +19,21 @@ dependencies = [ [project.urls] Homepage = "https://github.com/pgvector/pgvector-python" +[dependency-groups] +dev = [ + "asyncpg", + "Django", + "peewee", + "pg8000", + "psycopg[binary,pool]", + "psycopg2-binary", + "pytest", + "pytest-asyncio", + "scipy", + "SQLAlchemy[asyncio]>=2", + "sqlmodel>=0.0.12" +] + [tool.pytest.ini_options] asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index a13be06..0000000 --- a/requirements.txt +++ /dev/null @@ -1,12 +0,0 @@ -asyncpg -Django -numpy -peewee -pg8000 -psycopg[binary,pool] -psycopg2-binary -pytest -pytest-asyncio -scipy -SQLAlchemy[asyncio]>=2 -sqlmodel>=0.0.12 diff --git a/tests/test_bit.py b/tests/test_bit.py index 5e1bff2..5a71642 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -7,26 +7,42 @@ class TestBit: def test_list(self): assert Bit([True, False, True]).to_list() == [True, False, True] + def test_list_none(self): + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit([True, None, True]).to_text() == '101' + + def test_list_int(self): + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit([254, 7, 0]).to_text() == '110' + def test_tuple(self): assert Bit((True, False, True)).to_list() == [True, False, True] def test_str(self): assert Bit('101').to_list() == [True, False, True] + def test_bytes(self): + assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' + assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' + + def test_ndarray(self): + arr = np.array([True, False, True]) + assert Bit(arr).to_list() == [True, False, True] + assert np.array_equal(Bit(arr).to_numpy(), arr) + + def test_ndarray_unpackbits(self): + arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) + assert Bit(arr).to_text() == '111111100000011100000000' + def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) - assert Bit(arr).to_text() == '111111100000011100000000' + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit(arr).to_text() == '110' def test_ndarray_uint16(self): arr = np.array([254, 7, 0], dtype=np.uint16) - with pytest.raises(ValueError) as error: - Bit(arr) - assert str(error.value) == 'expected dtype to be bool or uint8' - - def test_ndarray_same_object(self): - arr = np.array([True, False, True]) - assert Bit(arr).to_list() == [True, False, True] - assert Bit(arr).to_numpy() is arr + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit(arr).to_text() == '110' def test_ndim_two(self): with pytest.raises(ValueError) as error: diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 4d3e474..61fbc4c 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -1,10 +1,10 @@ +from getpass import getuser import numpy as np -import os from pgvector import HalfVector, SparseVector, Vector from pgvector.pg8000 import register_vector from pg8000.native import Connection -conn = Connection(os.environ["USER"], database='pgvector_python_test') +conn = Connection(getuser(), database='pgvector_python_test') conn.run('CREATE EXTENSION IF NOT EXISTS vector') conn.run('DROP TABLE IF EXISTS pg8000_items') diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index dff03dd..d580f32 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,7 +1,7 @@ import numpy as np from pgvector import SparseVector import pytest -from scipy.sparse import coo_array +from scipy.sparse import coo_array, coo_matrix, csr_array, csr_matrix from struct import pack @@ -43,12 +43,30 @@ def test_coo_array_dimensions(self): SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) assert str(error.value) == 'extra argument' + def test_coo_matrix(self): + mat = coo_matrix(np.array([1, 0, 2, 0, 3, 0])) + vec = SparseVector(mat) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + def test_dok_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] + def test_csr_array(self): + arr = csr_array(np.array([[1, 0, 2, 0, 3, 0]])) + vec = SparseVector(arr) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + + def test_csr_matrix(self): + mat = csr_matrix(np.array([1, 0, 2, 0, 3, 0])) + vec = SparseVector(mat) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + def test_repr(self): assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 0d8d1ca..7558942 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,6 +1,6 @@ import asyncpg +from getpass import getuser import numpy as np -import os from pgvector import HalfVector, SparseVector, Vector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest @@ -28,7 +28,7 @@ def psycopg2_connect(dbapi_connection, connection_record): register_vector(dbapi_connection) -pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') +pg8000_engine = create_engine(f'postgresql+pg8000://{getuser()}@localhost/pgvector_python_test') if sqlalchemy_version > 1: psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') @@ -43,7 +43,7 @@ def psycopg_connect(dbapi_connection, connection_record): psycopg_async_type_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') @event.listens_for(psycopg_async_type_engine.sync_engine, "connect") - def connect(dbapi_connection, connection_record): + def psycopg_async_connect(dbapi_connection, connection_record): from pgvector.psycopg import register_vector_async dbapi_connection.run_async(register_vector_async) @@ -51,7 +51,7 @@ def connect(dbapi_connection, connection_record): asyncpg_type_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') @event.listens_for(asyncpg_type_engine.sync_engine, "connect") - def connect(dbapi_connection, connection_record): + def asyncpg_connect(dbapi_connection, connection_record): from pgvector.asyncpg import register_vector dbapi_connection.run_async(register_vector) @@ -103,7 +103,6 @@ class Item(Base): 'sqlalchemy_orm_half_precision_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'} ) half_precision_index.create(setup_engine) @@ -112,7 +111,6 @@ class Item(Base): 'sqlalchemy_orm_binary_quantize_index', func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'), postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'bit_hamming_ops'} ) binary_quantize_index.create(setup_engine) @@ -398,6 +396,20 @@ def test_sparsevec_l1_distance_orm(self, engine): items = session.scalars(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_subquery(self, engine): + create_items() + with Session(engine) as session: + subquery = select(Item.embedding).filter_by(id=1).scalar_subquery() + items = session.query(Item).order_by(Item.embedding.l2_distance(subquery)).all() + assert [v.id for v in items] == [1, 3, 2] + + def test_subquery_orm(self, engine): + create_items() + with Session(engine) as session: + subquery = select(Item.embedding).filter_by(id=1).scalar_subquery() + items = session.scalars(select(Item).order_by(Item.embedding.l2_distance(subquery))) + assert [v.id for v in items] == [1, 3, 2] + def test_filter(self, engine): create_items() with Session(engine) as session: @@ -528,6 +540,22 @@ def test_binary_quantize(self, engine): items = session.query(Item).order_by(distance).all() assert [v.id for v in items] == [2, 3, 1] + def test_binary_quantize_reranking(self, engine): + # recreate index (could also vacuum table) + binary_quantize_index.drop(setup_engine) + binary_quantize_index.create(setup_engine) + + with Session(engine) as session: + session.add(Item(id=1, embedding=[-1, -2, -3])) + session.add(Item(id=2, embedding=[1, -2, 3])) + session.add(Item(id=3, embedding=[1, 2, 3])) + session.commit() + + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) + subquery = session.query(Item).order_by(distance).limit(20).subquery() + items = session.query(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5).all() + assert [v.id for v in items] == [2, 3, 1] + @pytest.mark.parametrize('engine', array_engines) class TestSqlalchemyArray: @@ -596,6 +624,11 @@ async def test_bit(self, engine): item = await session.get(Item, 1) assert item.binary_embedding == embedding + if engine == asyncpg_engine: + session.add(Item(id=2, binary_embedding='101')) + item = await session.get(Item, 2) + assert item.binary_embedding == embedding + await engine.dispose() @pytest.mark.asyncio