Updated examples [skip ci]

Author: ankane

examples/lightfm/example.py

@@ -2,7 +2,7 @@
 from lightfm.datasets import fetch_movielens
 from pgvector.sqlalchemy import VECTOR
 from sqlalchemy import create_engine, insert, select, text, Float, Integer, String
-from sqlalchemy.orm import mapped_column, DeclarativeBase, Session
+from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session
 
 engine = create_engine('postgresql+psycopg://localhost/pgvector_example')
 with engine.connect() as conn:
@@ -17,17 +17,17 @@ class Base(DeclarativeBase):
 class User(Base):
     __tablename__ = 'user'
 
-    id = mapped_column(Integer, primary_key=True)
-    factors = mapped_column(VECTOR(20))
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    factors: Mapped[list[float]] = mapped_column(VECTOR(20))
 
 
 class Item(Base):
     __tablename__ = 'item'
 
-    id = mapped_column(Integer, primary_key=True)
-    title = mapped_column(String)
-    factors = mapped_column(VECTOR(20))
-    bias = mapped_column(Float)
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    title: Mapped[str] = mapped_column(String)
+    factors: Mapped[list[float]] = mapped_column(VECTOR(20))
+    bias: Mapped[float] = mapped_column(Float)
 
 
 Base.metadata.drop_all(engine)

examples/openai/halfvec.py

@@ -1,5 +1,6 @@
 from openai import OpenAI
-from pgvector.psycopg import register_vector, HalfVector
+from pgvector import HalfVector
+from pgvector.psycopg import register_vector
 import psycopg
 
 conn = psycopg.connect(dbname='pgvector_example', autocommit=True)

examples/surprise/example.py

@@ -1,6 +1,6 @@
 from pgvector.sqlalchemy import VECTOR
 from sqlalchemy import create_engine, insert, select, text, Integer
-from sqlalchemy.orm import mapped_column, DeclarativeBase, Session
+from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session
 from surprise import Dataset, SVD
 
 engine = create_engine('postgresql+psycopg://localhost/pgvector_example')
@@ -16,15 +16,15 @@ class Base(DeclarativeBase):
 class User(Base):
     __tablename__ = 'user'
 
-    id = mapped_column(Integer, primary_key=True)
-    factors = mapped_column(VECTOR(20))
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    factors: Mapped[list[float]] = mapped_column(VECTOR(20))
 
 
 class Item(Base):
     __tablename__ = 'item'
 
-    id = mapped_column(Integer, primary_key=True)
-    factors = mapped_column(VECTOR(20))
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    factors: Mapped[list[float]] = mapped_column(VECTOR(20))
 
 
 Base.metadata.drop_all(engine)