Skip to content

Commit c6b11aa

Browse files
committed
feat: add a full example for sqlmodel and fix error in README.md.
1 parent 6526319 commit c6b11aa

File tree

2 files changed

+40
-1
lines changed

2 files changed

+40
-1
lines changed

README.md

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -214,10 +214,13 @@ session.exec(text('CREATE EXTENSION IF NOT EXISTS vector'))
214214
Add a vector column
215215

216216
```python
217+
from typing import List, Optional
218+
217219
from pgvector.sqlalchemy import Vector
218-
from sqlalchemy import Column
220+
from sqlmodel import Column, Field, Session, SQLModel, create_engine, select
219221

220222
class Item(SQLModel, table=True):
223+
id: Optional[int] = Field(default=None, primary_key=True)
221224
embedding: List[float] = Field(sa_column=Column(Vector(3)))
222225
```
223226

@@ -237,6 +240,8 @@ session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(
237240

238241
Also supports `max_inner_product` and `cosine_distance`
239242

243+
See [examples/simple_sqlmodel_vector.py](examples/simple_sqlmodel_vector.py) for full code.
244+
240245
Get the distance
241246

242247
```python

examples/simple_sqlmodel_vector.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
"""
2+
A simple sqlmodel vector demo via pgvector.
3+
4+
For mac, if depdency missing or error, try `pip install pgvector-binary`
5+
"""
6+
7+
from typing import List, Optional
8+
9+
from pgvector.sqlalchemy import Vector
10+
from sqlmodel import Column, Field, Session, SQLModel, create_engine, select
11+
12+
13+
class Item(SQLModel, table=True):
14+
id: Optional[int] = Field(default=None, primary_key=True)
15+
16+
embedding: List[float] = Field(sa_column=Column(Vector(3)))
17+
18+
sqlite_url = f"postgresql://testuser:testuser@localhost:5432/testdb"
19+
20+
engine = create_engine(sqlite_url, echo=False)
21+
22+
SQLModel.metadata.create_all(engine)
23+
24+
with Session(engine) as session:
25+
item = Item(embedding=[1, 2, 3])
26+
session.add(item)
27+
session.commit()
28+
29+
res = session.exec(
30+
select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)
31+
)
32+
33+
for i in res:
34+
print(i)

0 commit comments

Comments
 (0)