sqlmodel_items.py

import os

from azure.identity import DefaultAzureCredential
from dotenv import load_dotenv
from pgvector.sqlalchemy import Vector
from sqlalchemy import Column
from sqlmodel import Field, Session, SQLModel, create_engine, func, select


class Item(SQLModel, table=True):
    id: int = Field(primary_key=True, sa_column_kwargs={"autoincrement": True})
    embedding: list[float] = Field(sa_column=Column(Vector(3)))


# Connect to the database
load_dotenv(".env", override=True)
POSTGRES_HOST = os.environ["POSTGRES_HOST"]
POSTGRES_USERNAME = os.environ["POSTGRES_USERNAME"]
POSTGRES_DATABASE = os.environ["POSTGRES_DATABASE"]

if POSTGRES_HOST.endswith(".database.azure.com"):
    print("Authenticating to Azure Database for PostgreSQL using Azure Identity...")
    azure_credential = DefaultAzureCredential()
    token = azure_credential.get_token("https://ossrdbms-aad.database.windows.net/.default")
    POSTGRES_PASSWORD = token.token
else:
    POSTGRES_PASSWORD = os.environ["POSTGRES_PASSWORD"]

DATABASE_URI = f"postgresql://{POSTGRES_USERNAME}:{POSTGRES_PASSWORD}@{POSTGRES_HOST}/{POSTGRES_DATABASE}"
# Specify SSL mode if needed
if POSTGRES_SSL := os.environ.get("POSTGRES_SSL"):
    DATABASE_URI += f"?sslmode={POSTGRES_SSL}"

engine = create_engine(DATABASE_URI, echo=False)

SQLModel.metadata.drop_all(engine)
SQLModel.metadata.create_all(engine)

with Session(engine) as session:
    session.add_all(
        [
            Item(embedding=[1, 2, 3]),
            Item(embedding=[-1, 1, 3]),
            Item(embedding=[0, -1, -2]),
        ]
    )

    # Find 2 closest vectors to [3, 1, 2]
    closest_items = session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(2))
    for item in closest_items:
        print(item.embedding)

    # Calculate distance between [3, 1, 2] and the first vector
    distance = session.exec(select(Item.embedding.l2_distance([3, 1, 2]))).first()
    print(distance)

    # Find vectors within distance 5 from [3, 1, 2]
    close_enough_items = session.exec(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5))
    for item in close_enough_items:
        print(item.embedding)

    # Calculate average of all vectors
    avg_embedding = session.exec(select(func.avg(Item.embedding))).first()
    print(avg_embedding)