import asyncpg
import numpy as np
import os
from pgvector import HalfVector, SparseVector, Vector
from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum
import pytest
from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY
from sqlalchemy.exc import StatementError
from sqlalchemy.ext.automap import automap_base
from sqlalchemy.orm import declarative_base, Session
from sqlalchemy.sql import func

try:
    from sqlalchemy.orm import mapped_column
    from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
    sqlalchemy_version = 2
except ImportError:
    mapped_column = Column
    sqlalchemy_version = 1

psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test')
psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test')


@event.listens_for(psycopg2_type_engine, "connect")
def psycopg2_connect(dbapi_connection, connection_record):
    from pgvector.psycopg2 import register_vector
    register_vector(dbapi_connection)


pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test')

if sqlalchemy_version > 1:
    psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test')
    psycopg_type_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test')

    @event.listens_for(psycopg_type_engine, "connect")
    def psycopg_connect(dbapi_connection, connection_record):
        from pgvector.psycopg import register_vector
        register_vector(dbapi_connection)

    psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test')
    psycopg_async_type_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test')

    @event.listens_for(psycopg_async_type_engine.sync_engine, "connect")
    def connect(dbapi_connection, connection_record):
        from pgvector.psycopg import register_vector_async
        dbapi_connection.run_async(register_vector_async)

    asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test')
    asyncpg_type_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test')

    @event.listens_for(asyncpg_type_engine.sync_engine, "connect")
    def connect(dbapi_connection, connection_record):
        from pgvector.asyncpg import register_vector
        dbapi_connection.run_async(register_vector)

engines = [psycopg2_engine, psycopg2_type_engine, pg8000_engine]
array_engines = [psycopg2_type_engine]
async_engines = []
async_array_engines = []

if sqlalchemy_version > 1:
    engines += [psycopg_engine, psycopg_type_engine]
    array_engines += [psycopg_type_engine]
    # TODO support asyncpg_type_engine
    async_engines += [psycopg_async_engine, psycopg_async_type_engine, asyncpg_engine]
    async_array_engines += [psycopg_async_type_engine, asyncpg_engine]

setup_engine = engines[0]
with Session(setup_engine) as session:
    session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
    session.commit()

Base = declarative_base()


class Item(Base):
    __tablename__ = 'sqlalchemy_orm_item'

    id = mapped_column(Integer, primary_key=True)
    embedding = mapped_column(VECTOR(3))
    half_embedding = mapped_column(HALFVEC(3))
    binary_embedding = mapped_column(BIT(3))
    sparse_embedding = mapped_column(SPARSEVEC(3))
    embeddings = mapped_column(ARRAY(VECTOR(3)))
    half_embeddings = mapped_column(ARRAY(HALFVEC(3)))


Base.metadata.drop_all(setup_engine)
Base.metadata.create_all(setup_engine)

index = Index(
    'sqlalchemy_orm_index',
    Item.embedding,
    postgresql_using='hnsw',
    postgresql_with={'m': 16, 'ef_construction': 64},
    postgresql_ops={'embedding': 'vector_l2_ops'}
)
index.create(setup_engine)

half_precision_index = Index(
    'sqlalchemy_orm_half_precision_index',
    func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
    postgresql_using='hnsw',
    postgresql_with={'m': 16, 'ef_construction': 64},
    postgresql_ops={'embedding': 'halfvec_l2_ops'}
)
half_precision_index.create(setup_engine)

binary_quantize_index = Index(
    'sqlalchemy_orm_binary_quantize_index',
    func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'),
    postgresql_using='hnsw',
    postgresql_with={'m': 16, 'ef_construction': 64},
    postgresql_ops={'embedding': 'bit_hamming_ops'}
)
binary_quantize_index.create(setup_engine)


def create_items():
    with Session(setup_engine) as session:
        session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])))
        session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])))
        session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])))
        session.commit()


def delete_items():
    with Session(setup_engine) as session:
        session.query(Item).delete()
        session.commit()


@pytest.mark.parametrize('engine', engines)
class TestSqlalchemy:
    def setup_method(self):
        delete_items()

    def test_core(self, engine):
        metadata = MetaData()

        item_table = Table(
            'sqlalchemy_core_item',
            metadata,
            Column('id', Integer, primary_key=True),
            Column('embedding', VECTOR(3)),
            Column('half_embedding', HALFVEC(3)),
            Column('binary_embedding', BIT(3)),
            Column('sparse_embedding', SPARSEVEC(3)),
            Column('embeddings', ARRAY(VECTOR(3)))
        )

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

        ivfflat_index = Index(
            'sqlalchemy_core_ivfflat_index',
            item_table.c.embedding,
            postgresql_using='ivfflat',
            postgresql_with={'lists': 1},
            postgresql_ops={'embedding': 'vector_l2_ops'}
        )
        ivfflat_index.create(engine)

        hnsw_index = Index(
            'sqlalchemy_core_hnsw_index',
            item_table.c.embedding,
            postgresql_using='hnsw',
            postgresql_with={'m': 16, 'ef_construction': 64},
            postgresql_ops={'embedding': 'vector_l2_ops'}
        )
        hnsw_index.create(engine)

    def test_orm(self, engine):
        item = Item(embedding=np.array([1.5, 2, 3]))
        item2 = Item(embedding=[4, 5, 6])
        item3 = Item()

        with Session(engine) as session:
            session.add(item)
            session.add(item2)
            session.add(item3)
            session.commit()

        stmt = select(Item)
        with Session(engine) as session:
            items = [v[0] for v in session.execute(stmt).all()]
            # TODO improve
            assert items[0].id % 3 == 1
            assert items[1].id % 3 == 2
            assert items[2].id % 3 == 0
            assert items[0].embedding == Vector([1.5, 2, 3])
            assert items[1].embedding == Vector([4, 5, 6])
            assert items[2].embedding is None

    def test_vector(self, engine):
        with Session(engine) as session:
            session.add(Item(id=1, embedding=[1, 2, 3]))
            session.commit()
            item = session.get(Item, 1)
            assert item.embedding == Vector([1, 2, 3])

    def test_vector_l2_distance(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.embedding.l2_distance([1, 1, 1])).all()
            assert [v.id for v in items] == [1, 3, 2]

    def test_vector_l2_distance_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1])))
            assert [v.id for v in items] == [1, 3, 2]

    def test_vector_max_inner_product(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])).all()
            assert [v.id for v in items] == [2, 3, 1]

    def test_vector_max_inner_product_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])))
            assert [v.id for v in items] == [2, 3, 1]

    def test_vector_cosine_distance(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])).all()
            assert [v.id for v in items] == [1, 2, 3]

    def test_vector_cosine_distance_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])))
            assert [v.id for v in items] == [1, 2, 3]

    def test_vector_l1_distance(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 1])).all()
            assert [v.id for v in items] == [1, 3, 2]

    def test_vector_l1_distance_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1])))
            assert [v.id for v in items] == [1, 3, 2]

    def test_halfvec(self, engine):
        with Session(engine) as session:
            session.add(Item(id=1, half_embedding=[1, 2, 3]))
            session.commit()
            item = session.get(Item, 1)
            assert item.half_embedding == HalfVector([1, 2, 3])

    def test_halfvec_l2_distance(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1])).all()
            assert [v.id for v in items] == [1, 3, 2]

    def test_halfvec_l2_distance_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1])))
            assert [v.id for v in items] == [1, 3, 2]

    def test_halfvec_max_inner_product(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1])).all()
            assert [v.id for v in items] == [2, 3, 1]

    def test_halfvec_max_inner_product_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1])))
            assert [v.id for v in items] == [2, 3, 1]

    def test_halfvec_cosine_distance(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1])).all()
            assert [v.id for v in items] == [1, 2, 3]

    def test_halfvec_cosine_distance_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1])))
            assert [v.id for v in items] == [1, 2, 3]

    def test_halfvec_l1_distance(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1])).all()
            assert [v.id for v in items] == [1, 3, 2]

    def test_halfvec_l1_distance_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1])))
            assert [v.id for v in items] == [1, 3, 2]

    def test_bit(self, engine):
        with Session(engine) as session:
            session.add(Item(id=1, binary_embedding='101'))
            session.commit()
            item = session.get(Item, 1)
            assert item.binary_embedding == '101'

    def test_bit_hamming_distance(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all()
            assert [v.id for v in items] == [2, 3, 1]

    def test_bit_hamming_distance_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.binary_embedding.hamming_distance('101')))
            assert [v.id for v in items] == [2, 3, 1]

    def test_bit_jaccard_distance(self, engine):
        if engine == pg8000_engine:
            return

        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all()
            assert [v.id for v in items] == [2, 3, 1]

    def test_bit_jaccard_distance_orm(self, engine):
        if engine == pg8000_engine:
            return

        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.binary_embedding.jaccard_distance('101')))
            assert [v.id for v in items] == [2, 3, 1]

    def test_sparsevec(self, engine):
        with Session(engine) as session:
            session.add(Item(id=1, sparse_embedding=[1, 2, 3]))
            session.commit()
            item = session.get(Item, 1)
            assert item.sparse_embedding == SparseVector([1, 2, 3])

    def test_sparsevec_l2_distance(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1])).all()
            assert [v.id for v in items] == [1, 3, 2]

    def test_sparsevec_l2_distance_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1])))
            assert [v.id for v in items] == [1, 3, 2]

    def test_sparsevec_max_inner_product(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1])).all()
            assert [v.id for v in items] == [2, 3, 1]

    def test_sparsevec_max_inner_product_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1])))
            assert [v.id for v in items] == [2, 3, 1]

    def test_sparsevec_cosine_distance(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1])).all()
            assert [v.id for v in items] == [1, 2, 3]

    def test_sparsevec_cosine_distance_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1])))
            assert [v.id for v in items] == [1, 2, 3]

    def test_sparsevec_l1_distance(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1])).all()
            assert [v.id for v in items] == [1, 3, 2]

    def test_sparsevec_l1_distance_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1])))
            assert [v.id for v in items] == [1, 3, 2]

    def test_filter(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1).all()
            assert [v.id for v in items] == [1]

    def test_filter_orm(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.scalars(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1))
            assert [v.id for v in items] == [1]

    def test_select(self, engine):
        with Session(engine) as session:
            session.add(Item(embedding=[2, 3, 3]))
            items = session.query(Item.embedding.l2_distance([1, 1, 1])).first()
            assert items[0] == 3

    def test_select_orm(self, engine):
        with Session(engine) as session:
            session.add(Item(embedding=[2, 3, 3]))
            items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all()
            assert items[0] == 3

    def test_avg(self, engine):
        with Session(engine) as session:
            res = session.query(avg(Item.embedding)).first()[0]
            assert res is None
            session.add(Item(embedding=[1, 2, 3]))
            session.add(Item(embedding=[4, 5, 6]))
            res = session.query(avg(Item.embedding)).first()[0]
            assert res == Vector([2.5, 3.5, 4.5])

    def test_avg_orm(self, engine):
        with Session(engine) as session:
            res = session.scalars(select(avg(Item.embedding))).first()
            assert res is None
            session.add(Item(embedding=[1, 2, 3]))
            session.add(Item(embedding=[4, 5, 6]))
            res = session.scalars(select(avg(Item.embedding))).first()
            assert res == Vector([2.5, 3.5, 4.5])

    def test_sum(self, engine):
        with Session(engine) as session:
            res = session.query(sum(Item.embedding)).first()[0]
            assert res is None
            session.add(Item(embedding=[1, 2, 3]))
            session.add(Item(embedding=[4, 5, 6]))
            res = session.query(sum(Item.embedding)).first()[0]
            assert res == Vector([5, 7, 9])

    def test_sum_orm(self, engine):
        with Session(engine) as session:
            res = session.scalars(select(sum(Item.embedding))).first()
            assert res is None
            session.add(Item(embedding=[1, 2, 3]))
            session.add(Item(embedding=[4, 5, 6]))
            res = session.scalars(select(sum(Item.embedding))).first()
            assert res == Vector([5, 7, 9])

    def test_bad_dimensions(self, engine):
        item = Item(embedding=[1, 2])
        with Session(engine) as session:
            session.add(item)
            with pytest.raises(StatementError, match='expected 3 dimensions, not 2'):
                session.commit()

    def test_bad_ndim(self, engine):
        item = Item(embedding=np.array([[1, 2, 3]]))
        with Session(engine) as session:
            session.add(item)
            with pytest.raises(StatementError, match='expected ndim to be 1'):
                session.commit()

    def test_bad_dtype(self, engine):
        item = Item(embedding=np.array(['one', 'two', 'three']))
        with Session(engine) as session:
            session.add(item)
            with pytest.raises(StatementError, match='could not convert string to float'):
                session.commit()

    def test_inspect(self, engine):
        columns = inspect(engine).get_columns('sqlalchemy_orm_item')
        assert isinstance(columns[1]['type'], VECTOR)

    def test_literal_binds(self, engine):
        sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(engine, compile_kwargs={'literal_binds': True})
        assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql)

    def test_insert(self, engine):
        with Session(engine) as session:
            session.execute(insert(Item).values(embedding=np.array([1, 2, 3])))

    def test_insert_bulk(self, engine):
        with Session(engine) as session:
            session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}])

    # register_vector in psycopg2 tests change this behavior
    # def test_insert_text(self):
    #     with Session(engine) as session:
    #         session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])})

    def test_automap(self, engine):
        metadata = MetaData()
        metadata.reflect(engine, only=['sqlalchemy_orm_item'])
        AutoBase = automap_base(metadata=metadata)
        AutoBase.prepare()
        AutoItem = AutoBase.classes.sqlalchemy_orm_item
        with Session(engine) as session:
            session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}])
            item = session.query(AutoItem).first()
            assert item.embedding == Vector([1, 2, 3])

    def test_half_precision(self, engine):
        create_items()
        with Session(engine) as session:
            items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all()
            assert [v.id for v in items] == [1, 3, 2]

    def test_binary_quantize(self, engine):
        with Session(engine) as session:
            session.add(Item(id=1, embedding=[-1, -2, -3]))
            session.add(Item(id=2, embedding=[1, -2, 3]))
            session.add(Item(id=3, embedding=[1, 2, 3]))
            session.commit()

            distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
            items = session.query(Item).order_by(distance).all()
            assert [v.id for v in items] == [2, 3, 1]


@pytest.mark.parametrize('engine', array_engines)
class TestSqlalchemyArray:
    def setup_method(self):
        delete_items()

    def test_vector_array(self, engine):
        with Session(engine) as session:
            session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]))
            session.commit()

            # this fails if the driver does not cast arrays
            item = session.get(Item, 1)
            assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])]

    def test_halfvec_array(self, engine):
        with Session(engine) as session:
            session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]))
            session.commit()

            # this fails if the driver does not cast arrays
            item = session.get(Item, 1)
            assert item.half_embeddings[0] == HalfVector([1, 2, 3])
            assert item.half_embeddings[1] == HalfVector([4, 5, 6])


@pytest.mark.parametrize('engine', async_engines)
class TestSqlalchemyAsync:
    def setup_method(self):
        delete_items()

    @pytest.mark.asyncio
    async def test_vector(self, engine):
        async_session = async_sessionmaker(engine, expire_on_commit=False)

        async with async_session() as session:
            async with session.begin():
                embedding = Vector([1, 2, 3])
                session.add(Item(id=1, embedding=embedding))
                item = await session.get(Item, 1)
                assert item.embedding == embedding

        await engine.dispose()

    @pytest.mark.asyncio
    async def test_halfvec(self, engine):
        async_session = async_sessionmaker(engine, expire_on_commit=False)

        async with async_session() as session:
            async with session.begin():
                embedding = HalfVector([1, 2, 3])
                session.add(Item(id=1, half_embedding=embedding))
                item = await session.get(Item, 1)
                assert item.half_embedding == embedding

        await engine.dispose()

    @pytest.mark.asyncio
    async def test_bit(self, engine):
        async_session = async_sessionmaker(engine, expire_on_commit=False)

        async with async_session() as session:
            async with session.begin():
                embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101'
                session.add(Item(id=1, binary_embedding=embedding))
                item = await session.get(Item, 1)
                assert item.binary_embedding == embedding

        await engine.dispose()

    @pytest.mark.asyncio
    async def test_sparsevec(self, engine):
        async_session = async_sessionmaker(engine, expire_on_commit=False)

        async with async_session() as session:
            async with session.begin():
                embedding = SparseVector([1, 2, 3])
                session.add(Item(id=1, sparse_embedding=embedding))
                item = await session.get(Item, 1)
                assert item.sparse_embedding == embedding

        await engine.dispose()

    @pytest.mark.asyncio
    async def test_avg(self, engine):
        async_session = async_sessionmaker(engine, expire_on_commit=False)

        async with async_session() as session:
            async with session.begin():
                session.add(Item(embedding=[1, 2, 3]))
                session.add(Item(embedding=[4, 5, 6]))
                res = await session.scalars(select(avg(Item.embedding)))
                assert res.first() == Vector([2.5, 3.5, 4.5])

        await engine.dispose()


@pytest.mark.parametrize('engine', async_array_engines)
class TestSqlalchemyAsyncArray:
    def setup_method(self):
        delete_items()

    @pytest.mark.asyncio
    async def test_vector_array(self, engine):
        async_session = async_sessionmaker(engine, expire_on_commit=False)

        async with async_session() as session:
            async with session.begin():
                session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])]))
                item = await session.get(Item, 1)
                assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])]

                session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]))
                item = await session.get(Item, 2)
                assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])]

        await engine.dispose()