refactor: Add document loaders and update constants

Signed-off-by: teleprint-me <[email protected]>
mashimi · Jun 4, 2023 · 672e34f · 672e34f
1 parent 3d1806d
commit 672e34f
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/constants.py b/constants.py
@@ -1,6 +1,8 @@
 import os
+
 # from dotenv import load_dotenv
 from chromadb.config import Settings
+from langchain.document_loaders import CSVLoader, PDFMinerLoader, TextLoader
 
 # load_dotenv()
 ROOT_DIRECTORY = os.path.dirname(os.path.realpath(__file__))
@@ -15,4 +17,10 @@
         chroma_db_impl='duckdb+parquet',
         persist_directory=PERSIST_DIRECTORY,
         anonymized_telemetry=False
-)
+)
+
+DOCUMENT_MAP = {
+    ".txt": TextLoader,
+    ".pdf": PDFMinerLoader,
+    ".csv": CSVLoader,
+}