Skip to content

Commit a3e08bb

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

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
@@ -186,10 +186,13 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc
186186
Add a vector column
187187

188188
```python
189+
from typing import List, Optional
190+
189191
from pgvector.sqlalchemy import Vector
190-
from sqlalchemy import Column
192+
from sqlmodel import Column, Field, Session, SQLModel, create_engine, select
191193

192194
class Item(SQLModel, table=True):
195+
id: Optional[int] = Field(default=None, primary_key=True)
193196
embedding: List[float] = Field(sa_column=Column(Vector(3)))
194197
```
195198

@@ -209,6 +212,8 @@ session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(
209212

210213
Also supports `max_inner_product` and `cosine_distance`
211214

215+
See [examples/simple_sqlmodel_vector.py](examples/simple_sqlmodel_vector.py) for full code.
216+
212217
## Psycopg 3
213218

214219
Register the vector type with your connection

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)