Merge pull request #32 from 2405-team3/knowledge-base-refactor

Knowledge base refactor

.gitignore

@@ -130,6 +130,7 @@ venv/
 ENV/
 env.bak/
 venv.bak/
+.env.testing
 
 # Spyder project settings
 .spyderproject

Pipfile

@@ -28,11 +28,10 @@ motor = "*"
 rank-bm25 = "*"
 llama-index-postprocessor-colbert-rerank = "*"
 llama-index-retrievers-bm25 = "*"
-llama-index-storage-index-store-mongodb = "*"
-pylint = "*"
 
 [dev-packages]
 pytest = "*"
+pylint = "*"
 
 [requires]
 python_version = "3.10"

cleanup.py

@@ -0,0 +1,32 @@
+import os
+
+import pymongo
+from dotenv import load_dotenv
+
+load_dotenv(override=True, dotenv_path=".env.testing")
+
+MONGO_URI = os.environ["MONGO_URI"]
+CONFIG_DB = os.environ["CONFIG_DB"]
+CONFIG_KB_COL = os.environ["CONFIG_KB_COL"]
+
+kb_names = ["Sentence", "Semantic", "Markdown"]
+def drop_all_knowledge_bases():
+    print(f"Dropping all knowledge bases in {CONFIG_DB}")
+    client = pymongo.MongoClient(MONGO_URI)
+    client.drop_database(CONFIG_DB)
+    client.close()
+
+
+def remove_kb_files(kb_names: list):
+    client = pymongo.MongoClient(MONGO_URI)
+
+    for kb_name in kb_names:
+        client.drop_database(kb_name)
+
+    client.close()
+
+drop_all_knowledge_bases()
+remove_kb_files(kb_names)
+
+
+