From 884d7a173ecf6070781d29c8ceaf10a2904c1fd1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacy=20=C5=81=C4=85tka?= <74246391+LatekVo@users.noreply.github.com> Date: Mon, 15 Apr 2024 22:20:19 +0200 Subject: [PATCH] Fix allow_dangerous_deserialization bug (#18) Co-authored-by: Nekxis --- core/tools/dbops.py | 30 ++++++++++++++++++++++++------ 1 file changed, 24 insertions(+), 6 deletions(-) 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