diviner/service/data/db.py

# -*- coding: utf-8 -*-
from .__load__ import *
from langchain_community.vectorstores import FAISS

class Db(object):

    def __init__(self):
        Demeter.service('init')
        self.embedding = Demeter.embedding
        self.path = File.path() + 'data/db/'

    def write(self, id, data):
        name = self.name(id)
        if File.exists(name):
            db = FAISS.load_local(name, embeddings=self.embedding, allow_dangerous_deserialization=True)
            db.add_documents(data)
        else:
            db = FAISS.from_documents(data, embedding=self.embedding)
        db.save_local(name)
        return db

    def read(self, id):
        db = FAISS.load_local(self.name(id), embeddings=self.embedding, allow_dangerous_deserialization=True)
        return db

    def name(self, name):
        name = str(name)
        return File.dest(self.path, Demeter.md5(name))