pgvector 0.4.0__tar.gz → 0.4.2__tar.gz

{pgvector-0.4.0/pgvector.egg-info → pgvector-0.4.2}/PKG-INFO

@@ -1,14 +1,15 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: pgvector
-Version: 0.4.0
+Version: 0.4.2
 Summary: pgvector support for Python
 Author-email: Andrew Kane <andrew@ankane.org>
-License: MIT
+License-Expression: MIT
 Project-URL: Homepage, https://github.com/pgvector/pgvector-python
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
 Requires-Dist: numpy
+Dynamic: license-file
 
 # pgvector-python
 
@@ -39,6 +40,7 @@ And follow the instructions for your database library:
 
 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
@@ -270,7 +272,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'}
 )
 ```
@@ -282,6 +283,37 @@ 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
@@ -420,7 +452,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
@@ -483,7 +515,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
@@ -529,7 +561,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
@@ -583,7 +615,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

pgvector-0.4.0/PKG-INFO → pgvector-0.4.2/README.md

@@ -1,15 +1,3 @@
-Metadata-Version: 2.2
-Name: pgvector
-Version: 0.4.0
-Summary: pgvector support for Python
-Author-email: Andrew Kane <andrew@ankane.org>
-License: MIT
-Project-URL: Homepage, https://github.com/pgvector/pgvector-python
-Requires-Python: >=3.9
-Description-Content-Type: text/markdown
-License-File: LICENSE.txt
-Requires-Dist: numpy
-
 # pgvector-python
 
 [pgvector](https://github.com/pgvector/pgvector) support for Python
@@ -39,6 +27,7 @@ And follow the instructions for your database library:
 
 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
@@ -270,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'}
 )
 ```
@@ -282,6 +270,37 @@ 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
@@ -420,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
@@ -483,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
@@ -529,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
@@ -583,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

pgvector-0.4.2/pgvector/django/extensions.py

@@ -0,0 +1,11 @@
+from django import VERSION
+from django.contrib.postgres.operations import CreateExtension
+
+
+class VectorExtension(CreateExtension):
+    if VERSION[0] >= 6:
+        def __init__(self, hints=None):
+            super().__init__('vector', hints=hints)
+    else:
+        def __init__(self):
+            self.name = 'vector'

{pgvector-0.4.0 → pgvector-0.4.2}/pgvector/sparsevec.py

@@ -85,7 +85,7 @@ class SparseVector:
 
         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()

{pgvector-0.4.0 → pgvector-0.4.2}/pgvector/sqlalchemy/bit.py

@@ -14,6 +14,18 @@ class BIT(UserDefinedType):
             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)

pgvector-0.4.0/README.md → pgvector-0.4.2/pgvector.egg-info/PKG-INFO

@@ -1,3 +1,16 @@
+Metadata-Version: 2.4
+Name: pgvector
+Version: 0.4.2
+Summary: pgvector support for Python
+Author-email: Andrew Kane <andrew@ankane.org>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/pgvector/pgvector-python
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE.txt
+Requires-Dist: numpy
+Dynamic: license-file
+
 # pgvector-python
 
 [pgvector](https://github.com/pgvector/pgvector) support for Python
@@ -27,6 +40,7 @@ And follow the instructions for your database library:
 
 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
@@ -258,7 +272,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,6 +283,37 @@ 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
@@ -408,7 +452,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 +515,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 +561,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 +615,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

{pgvector-0.4.0 → pgvector-0.4.2}/pyproject.toml

@@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pgvector"
-version = "0.4.0"
+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"

{pgvector-0.4.0 → pgvector-0.4.2}/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 @@ class TestSparseVector:
             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)'

{pgvector-0.4.0 → pgvector-0.4.2}/tests/test_sqlalchemy.py

@@ -43,7 +43,7 @@ if sqlalchemy_version > 1:
     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 @@ if sqlalchemy_version > 1:
     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 @@ half_precision_index = Index(
     '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 @@ binary_quantize_index = Index(
     '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)
@@ -528,6 +526,22 @@ class TestSqlalchemy:
             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 +610,11 @@ class TestSqlalchemyAsync:
                 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

pgvector-0.4.0/pgvector/django/extensions.py

@@ -1,6 +0,0 @@
-from django.contrib.postgres.operations import CreateExtension
-
-
-class VectorExtension(CreateExtension):
-    def __init__(self):
-        self.name = 'vector'