Added support for Peewee - closes pgvector#26 and closes pgvector#27

ankane · bormanjo · ankane · commit ff9012eb03ce · 2023-09-05T12:59:20.000-07:00
Co-authored-by: John-Craig Borman &lt;borman.johncraig@gmail.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,6 @@
 ## 0.2.2 (unreleased)
 
+- Added support for Peewee
 - Added `HnswIndex` for Django
 
 ## 0.2.1 (2023-07-31)
diff --git a/README.md b/README.md
@@ -2,7 +2,7 @@
 
 [pgvector](https://github.com/pgvector/pgvector) support for Python
 
-Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), and [asyncpg](https://github.com/MagicStack/asyncpg)
+Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), and [Peewee](https://github.com/coleifer/peewee)
 
 [![Build Status](https://github.com/pgvector/pgvector-python/workflows/build/badge.svg?branch=master)](https://github.com/pgvector/pgvector-python/actions)
 
@@ -22,6 +22,7 @@ And follow the instructions for your database library:
 - [Psycopg 3](#psycopg-3)
 - [Psycopg 2](#psycopg-2)
 - [asyncpg](#asyncpg)
+- [Peewee](#peewee) [unreleased]
 
 Or check out some examples:
 
@@ -270,6 +271,43 @@ Get the nearest neighbors to a vector
 await conn.fetch('SELECT * FROM item ORDER BY embedding <-> $1 LIMIT 5', embedding)
 ```
 
+## Peewee
+
+Add a vector column
+
+```python
+from pgvector.peewee import VectorField
+
+class Item(BaseModel):
+    embedding = VectorField(dimensions=3)
+```
+
+Insert a vector
+
+```python
+item = Item.create(embedding=[1, 2, 3])
+```
+
+Get the nearest neighbors to a vector
+
+```python
+Item.select().order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)
+```
+
+Also supports `max_inner_product` and `cosine_distance`
+
+Get the distance
+
+```python
+Item.select(Item.embedding.l2_distance([3, 1, 2]).alias('distance'))
+```
+
+Get items within a certain distance
+
+```python
+Item.select().where(Item.embedding.l2_distance([3, 1, 2]) < 5)
+```
+
 ## History
 
 View the [changelog](https://github.com/pgvector/pgvector-python/blob/master/CHANGELOG.md)
diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py
@@ -0,0 +1,31 @@
+from peewee import Expression, Field, Value
+from ..utils import from_db, to_db
+
+
+class VectorField(Field):
+    field_type = 'vector'
+
+    def __init__(self, dimensions=None, *args, **kwargs):
+        self.dimensions = dimensions
+        super(VectorField, self).__init__(*args, **kwargs)
+
+    def get_modifiers(self):
+        return self.dimensions and [self.dimensions] or None
+
+    def db_value(self, value):
+        return to_db(value)
+
+    def python_value(self, value):
+        return from_db(value)
+
+    def _distance(self, op, vector):
+        return Expression(lhs=self, op=op, rhs=Value(vector, converter=to_db, unpack=False))
+
+    def l2_distance(self, vector):
+        return self._distance('<->', vector)
+
+    def max_inner_product(self, vector):
+        return self._distance('<#>', vector)
+
+    def cosine_distance(self, vector):
+        return self._distance('<=>', vector)
diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,7 @@
 asyncpg
 Django
 numpy
+peewee
 psycopg[binary]
 psycopg2-binary
 pytest
diff --git a/setup.py b/setup.py
@@ -16,6 +16,7 @@
     packages=[
         'pgvector.asyncpg',
         'pgvector.django',
+        'pgvector.peewee',
         'pgvector.psycopg',
         'pgvector.psycopg2',
         'pgvector.sqlalchemy',
diff --git a/tests/test_peewee.py b/tests/test_peewee.py
@@ -0,0 +1,68 @@
+from math import sqrt
+import numpy as np
+from peewee import Model, PostgresqlDatabase
+from pgvector.peewee import VectorField
+
+db = PostgresqlDatabase('pgvector_python_test')
+
+
+class BaseModel(Model):
+    class Meta:
+        database = db
+
+
+class Item(BaseModel):
+    embedding = VectorField(dimensions=3)
+
+
+db.connect()
+db.execute_sql('CREATE EXTENSION IF NOT EXISTS vector')
+db.drop_tables([Item])
+db.create_tables([Item])
+
+
+def create_items():
+    vectors = [
+        [1, 1, 1],
+        [2, 2, 2],
+        [1, 1, 2]
+    ]
+    for i, v in enumerate(vectors):
+        Item.create(id=i + 1, embedding=v)
+
+
+class TestPeewee:
+    def setup_method(self, test_method):
+        Item.truncate_table()
+
+    def test_works(self):
+        Item.create(id=1, embedding=[1, 2, 3])
+        item = Item.get_by_id(1)
+        assert np.array_equal(item.embedding, np.array([1, 2, 3]))
+        assert item.embedding.dtype == np.float32
+
+    def test_l2_distance(self):
+        create_items()
+        distance = Item.embedding.l2_distance([1, 1, 1])
+        items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5)
+        assert [v.id for v in items] == [1, 3, 2]
+        assert [v.distance for v in items] == [0, 1, sqrt(3)]
+
+    def test_max_inner_product(self):
+        create_items()
+        distance = Item.embedding.max_inner_product([1, 1, 1])
+        items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5)
+        assert [v.id for v in items] == [2, 3, 1]
+        assert [v.distance for v in items] == [-6, -4, -3]
+
+    def test_cosine_distance(self):
+        create_items()
+        distance = Item.embedding.cosine_distance([1, 1, 1])
+        items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5)
+        assert [v.id for v in items] == [1, 2, 3]
+        assert [v.distance for v in items] == [0, 0, 0.05719095841793653]
+
+    def test_where(self):
+        create_items()
+        items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1)
+        assert [v.id for v in items] == [1]

-Original file line number
+Diff line change
@@ @@ -1,6 +1,7 @@ @@
 asyncpg
 Django
 numpy
 +peewee
 psycopg[binary]
 psycopg2-binary
 pytest