SQLAlchemy

이 예제는 AkasicDB vector column을 가진 ORM model을 정의하고, vector index를 만든 뒤 L2 distance 기준으로 row를 정렬합니다.

설치

pip install "akasicdb[sqlalchemy]"

예제

from sqlalchemy import Index, create_engine, insert, select
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column

from akasicdb.sqlalchemy import VECTOR
from akasicdb.types import Hnsw, IndexOption

DB_URL = "postgresql+psycopg://postgres:mysecretpassword@localhost:5432/postgres"


class Base(DeclarativeBase):
    pass


class Document(Base):
    __tablename__ = "documents"

    id: Mapped[int] = mapped_column(primary_key=True)
    embedding = mapped_column(VECTOR(3))


engine = create_engine(DB_URL)
Base.metadata.drop_all(engine)
Base.metadata.create_all(engine)

with Session(engine) as session:
    session.execute(insert(Document).values(embedding=[1.0, 2.0, 3.0]))
    session.execute(insert(Document).values(embedding=[1.0, 3.0, 4.0]))
    session.commit()

    index = Index(
        "documents_embedding_idx",
        Document.embedding,
        postgresql_using="vectoron",
        postgresql_with={
            "options": f"$${IndexOption(index=Hnsw(), threads=1).dumps()}$$"
        },
        postgresql_ops={"embedding": "vector_l2_ops"},
    )
    index.create(session.bind)

    rows = session.scalars(
        select(Document).order_by(Document.embedding.l2_distance([1.0, 2.0, 3.5]))
    ).all()

    for row in rows:
        print(row.id, row.embedding)