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")