vec_txt_table.py

from .k_db import PostgresDB

TABLE_VEC_TXT = """
CREATE TABLE vec_txt (
    vector_id varchar(36) PRIMARY KEY, 
    text text, 
    paragraph_id varchar(40) not null
)
"""


# 025a9bee-2eb2-47f5-9722-525e05a0442b
class TxtVector:
    def __init__(self, db: PostgresDB) -> None:
        self.db = db

    def insert(self, vectors):
        query = f"INSERT INTO vec_txt(vector_id,text,paragraph_id) VALUES"
        args = []
        for value in vectors:
            value = list(value)
            query += "(%s,%s,%s),"
            args.extend(value)
        query = query[:len(query) - 1]
        query += f"ON conflict(vector_id) DO UPDATE SET text = EXCLUDED.text,paragraph_id = EXCLUDED.paragraph_id;"
        # query += ";"
        self.db.execute_args(query, args)

    def delete(self, ids):
        for item in ids:
            query = f"delete FROM vec_txt WHERE vector_id = '%s'" % (item,)
            self.db.execute(query)

    def search(self, search: str):
        query = f"SELECT paragraph_id,text FROM vec_txt WHERE vector_id = %s"
        self.db.execute_args(query, [search])
        answer = self.db.fetchall()
        # print(answer)
        return answer[0]

    def create_table(self):
        query = f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'vec_txt')"
        self.db.execute(query)
        exists = self.db.fetchall()[0][0]
        if not exists:
            query = TABLE_VEC_TXT
            self.db.execute(query)

    def drop_table(self):
        query = f"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'vec_txt')"
        self.db.execute(query)
        exists = self.db.fetchall()[0][0]
        if exists:
            query = "DROP TABLE vec_txt"
            self.db.format(query)
            print("drop table vec_txt ok")