pgvector_sqlalchemy_vector_column_nearest_neighbor_search.py

1from sqlalchemy import create_engine, Column, Integer
2from sqlalchemy.orm import declarative_base, Session
3from pgvector.sqlalchemy import Vector
4
5# Connect to the database
6engine = create_engine('postgresql://localhost/pgvector_example')
7
8# Define the model
9Base = declarative_base()
10
11class Item(Base):
12    __tablename__ = 'item'
13    id = Column(Integer, primary_key=True)
14    embedding = Column(Vector(3))
15
16# Create the table
17Base.metadata.drop_all(engine)
18Base.metadata.create_all(engine)
19
20# Insert data
21with Session(engine) as session:
22    item1 = Item(embedding=[1, 1, 1])
23    item2 = Item(embedding=[2, 2, 2])
24    item3 = Item(embedding=[1, 1, 2])
25    session.add_all([item1, item2, item3])
26    session.commit()
27
28# Perform a nearest neighbor search
29with Session(engine) as session:
30    items = session.query(Item).order_by(Item.embedding.l2_distance([1, 1, 1])).limit(5).all()
31    for item in items:
32        print(f'ID: {item.id}, Embedding: {item.embedding}')