diff --git a/core/tools/dbops.py b/core/tools/dbops.py index dc40724..3c25e86 100644 --- a/core/tools/dbops.py +++ b/core/tools/dbops.py @@ -5,14 +5,32 @@ def create_db_if_not_exists(db_name: str, embeddings: Embeddings): - if not exists('store/vector/' + db_name + '.faiss'): - print("Creating new database:", db_name + '.faiss') - tmp_db = FAISS.from_texts(['You are a large language model, intended for research purposes.'], embeddings) - tmp_db.save_local(folder_path='store/vector', index_name=db_name) + if not exists("store/vector/" + db_name + ".faiss"): + print("Creating new database:", db_name + ".faiss") + tmp_db = FAISS.from_texts( + ["You are a large language model, intended for research purposes."], + embeddings, + ) + tmp_db.save_local(folder_path="store/vector", index_name=db_name) else: - print("Already exists:", db_name + '.faiss') + print("Already exists:", db_name + ".faiss") def get_db_by_name(db_name: str, embeddings: Embeddings) -> FAISS: create_db_if_not_exists(db_name, embeddings) - return FAISS.load_local(folder_path='store/vector', embeddings=embeddings, index_name=db_name, allow_dangerous_deserialization=True) + + try: + # windows + db_connection = FAISS.load_local( + folder_path="store/vector", + embeddings=embeddings, + index_name=db_name, + allow_dangerous_deserialization=True, + ) + except Exception: + # linux & mac + db_connection = FAISS.load_local( + folder_path="store/vector", embeddings=embeddings, index_name=db_name + ) + + return db_connection