forked from pgvector/pgvector-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_sqlmodel.py
More file actions
129 lines (107 loc) · 4.4 KB
/
Copy pathtest_sqlmodel.py
File metadata and controls
129 lines (107 loc) · 4.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import numpy as np
from pgvector.sqlalchemy import Vector
import pytest
from sqlalchemy import Column, Index
from sqlalchemy.exc import StatementError
from sqlalchemy.sql import func
from sqlmodel import Field, Session, SQLModel, create_engine, delete, select, text
from typing import List, Optional
engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test')
with Session(engine) as session:
session.exec(text('CREATE EXTENSION IF NOT EXISTS vector'))
class Item(SQLModel, table=True):
__tablename__ = 'sqlmodel_item'
id: Optional[int] = Field(default=None, primary_key=True)
embedding: Optional[List[float]] = Field(default=None, sa_column=Column(Vector(3)))
SQLModel.metadata.drop_all(engine)
SQLModel.metadata.create_all(engine)
index = Index(
'sqlmodel_index',
Item.embedding,
postgresql_using='hnsw',
postgresql_with={'m': 16, 'ef_construction': 64},
postgresql_ops={'embedding': 'vector_l2_ops'}
)
index.create(engine)
def create_items():
vectors = [
[1, 1, 1],
[2, 2, 2],
[1, 1, 2]
]
session = Session(engine)
for i, v in enumerate(vectors):
session.add(Item(id=i + 1, embedding=v))
session.commit()
class TestSqlmodel:
def setup_method(self, test_method):
with Session(engine) as session:
session.exec(delete(Item))
session.commit()
def test_orm(self):
item = Item(embedding=[1.5, 2, 3])
item2 = Item(embedding=[4, 5, 6])
item3 = Item()
session = Session(engine)
session.add(item)
session.add(item2)
session.add(item3)
session.commit()
stmt = select(Item)
with Session(engine) as session:
items = session.exec(stmt).all()
assert items[0].id == 1
assert items[1].id == 2
assert items[2].id == 3
assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3]))
assert items[0].embedding.dtype == np.float32
assert np.array_equal(items[1].embedding, np.array([4, 5, 6]))
assert items[1].embedding.dtype == np.float32
assert items[2].embedding is None
def test_l2_distance(self):
create_items()
with Session(engine) as session:
items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1])))
assert [v.id for v in items] == [1, 3, 2]
def test_max_inner_product(self):
create_items()
with Session(engine) as session:
items = session.exec(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])))
assert [v.id for v in items] == [2, 3, 1]
def test_cosine_distance(self):
create_items()
with Session(engine) as session:
items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])))
assert [v.id for v in items] == [1, 2, 3]
def test_filter(self):
create_items()
with Session(engine) as session:
items = session.exec(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1))
assert [v.id for v in items] == [1]
def test_select(self):
with Session(engine) as session:
session.add(Item(embedding=[2, 3, 3]))
items = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all()
assert items[0] == 3
def test_avg(self):
with Session(engine) as session:
avg = session.exec(select(func.avg(Item.embedding))).first()
assert avg is None
session.add(Item(embedding=[1, 2, 3]))
session.add(Item(embedding=[4, 5, 6]))
avg = session.exec(select(func.avg(Item.embedding))).first()
assert np.array_equal(avg, np.array([2.5, 3.5, 4.5]))
def test_sum(self):
with Session(engine) as session:
sum = session.exec(select(func.sum(Item.embedding))).first()
assert sum is None
session.add(Item(embedding=[1, 2, 3]))
session.add(Item(embedding=[4, 5, 6]))
sum = session.exec(select(func.sum(Item.embedding))).first()
assert np.array_equal(sum, np.array([5, 7, 9]))
def test_bad_dimensions(self):
item = Item(embedding=[1, 2])
session = Session(engine)
session.add(item)
with pytest.raises(StatementError, match='expected 3 dimensions, not 2'):
session.commit()