isekai-toolkit/install.py

import asyncio
import asyncpg
import os
from config import Config

conn = None

class Install:
    dbi: asyncpg.Connection

    async def run(self):
        db_config = Config.get("database")
        self.dbi = await asyncpg.connect(db_config)
        args = os.sys.argv
        if "--force" in args:
            await self.drop_table()

        await self.create_table()

    async def drop_table(self):
        await self.dbi.execute("DROP TABLE IF EXISTS embedding_search_title_index;")
        print("Table dropped")

    async def create_table(self):
        embedding_vector_size = Config.get("chatcomplete.embedding_vector_size", 512, int)

        await self.dbi.execute("""
        CREATE TABLE embedding_search_title_index (
            id SERIAL PRIMARY KEY,
            sha1 VARCHAR(40) NOT NULL UNIQUE,
            title TEXT NOT NULL,
            rev_id INT8 NOT NULL,
            embedding VECTOR(%d) NOT NULL
        );
        """ % (embedding_vector_size))
        await self.dbi.execute("CREATE INDEX embedding_search_title_index_embedding_idx ON embedding_search_title_index USING ivfflat (embedding vector_cosine_ops);")
        print("Table created")

if __name__ == "__main__":
    install = Install()

    loop = asyncio.get_event_loop()
    loop.run_until_complete(install.run())