Update sqlalchemy_items.py

Author: nachoalonsoportillo

examples/sqlalchemy_items.py

@@ -39,23 +39,30 @@ class Item(Base):
 
 engine = create_engine(DATABASE_URI, echo=False)
 
-# Create tables in database
+# Create pgvector extension
+with engine.begin() as conn:
+ conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
+
+# Drop all tables defined in this model from the database, if they already exist
 Base.metadata.drop_all(engine)
+# Create all tables defined in this model in the database
 Base.metadata.create_all(engine)
 
 # Insert data and issue queries
 with Session(engine) as session:
- session.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
-
+ # Define HNSW index to support vector similarity search through the vector_l2_ops access method (Euclidean distance). The SQL operator for Euclidean distance is written as <->.
  index = Index(
- "my_index",
+ "hnsw_index_for_euclidean_distance_similarity_search",
  Item.embedding,
  postgresql_using="hnsw",
  postgresql_with={"m": 16, "ef_construction": 64},
  postgresql_ops={"embedding": "vector_l2_ops"},
  )
+
+ # Create the HNSW index
  index.create(engine)
 
+ # Insert three vectors as three separate rows in the items table
  session.add_all(
  [
  Item(embedding=[1, 2, 3]),
@@ -64,20 +71,28 @@ class Item(Base):
  ]
  )
 
+ # Find all vectors in table items
+ all_items = session.scalars(select(Item))
+ print("All vectors in table items:")
+ for item in all_items:
+ print(f"\t{item.embedding}")
+
  # Find 2 closest vectors to [3, 1, 2]
  closest_items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(2))
+ print("Two closest vectors to [3, 1, 2] in table items:")
  for item in closest_items:
- print(item.embedding)
+ print(f"\t{item.embedding}")
 
  # Calculate distance between [3, 1, 2] and the first vector
  distance = session.scalars(select(Item.embedding.l2_distance([3, 1, 2]))).first()
- print(distance)
+ print(f"Distance between [3, 1, 2] vector and the one closest to it: {distance}")
 
  # Find vectors within distance 5 from [3, 1, 2]
  close_enough_items = session.scalars(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5))
+ print("Vectors within a distance of 5 from [3, 1, 2]:")
  for item in close_enough_items:
- print(item.embedding)
+ print(f"\t{item.embedding}")
 
  # Calculate average of all vectors
  avg_embedding = session.scalars(select(func.avg(Item.embedding))).first()
- print(avg_embedding)
+ print(f"Average of all vectors: {avg_embedding}")