rnekrasov/pgvector-python: pgvector support for Python

tests/test_asyncpg.py
54 строки2 KB
Andrew Kane
Added pool example and test for asyncpg - closes #18
29 июн 2023, 01:30
29 июн 2023, 01:3034db96b
import asyncio
import asyncpg
import numpy as np
from pgvector.asyncpg import register_vector
import pytest


class TestAsyncpg:
    @pytest.mark.asyncio
    async def test_works(self):
        conn = await asyncpg.connect(database='pgvector_python_test')
        await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
        await conn.execute('DROP TABLE IF EXISTS item')
        await conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))')

        await register_vector(conn)

        embedding = np.array([1.5, 2, 3])
        await conn.execute("INSERT INTO item (embedding) VALUES ($1), (NULL)", embedding)

        res = await conn.fetch("SELECT * FROM item ORDER BY id")
        assert res[0]['id'] == 1
        assert res[1]['id'] == 2
        assert np.array_equal(res[0]['embedding'], embedding)
        assert res[0]['embedding'].dtype == np.float32
        assert res[1]['embedding'] is None

        # ensures binary format is correct
        text_res = await conn.fetch("SELECT embedding::text FROM item ORDER BY id LIMIT 1")
        assert text_res[0]['embedding'] == '[1.5,2,3]'

        await conn.close()

    @pytest.mark.asyncio
    async def test_pool(self):
        async def init(conn):
            await register_vector(conn)

        pool = await asyncpg.create_pool(database='pgvector_python_test', init=init)

        async with pool.acquire() as conn:
            await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
            await conn.execute('DROP TABLE IF EXISTS item')
            await conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))')

            embedding = np.array([1.5, 2, 3])
            await conn.execute("INSERT INTO item (embedding) VALUES ($1), (NULL)", embedding)

            res = await conn.fetch("SELECT * FROM item ORDER BY id")
            assert res[0]['id'] == 1
            assert res[1]['id'] == 2
            assert np.array_equal(res[0]['embedding'], embedding)
            assert res[0]['embedding'].dtype == np.float32
            assert res[1]['embedding'] is None