Update app.py

app.py

@@ -96,7 +96,7 @@ def load_model():
     return llm
     
 
-def create_vector_database():
+def create_vector_database(uploaded_files):
     # DB_DIR: str = os.path.join(ABS_PATH, "db")
     """
     Creates a vector database using document loaders and embeddings.
@@ -107,27 +107,50 @@ def create_vector_database():
 
     """
     # Initialize loaders for different file types
-    pdf_loader = DirectoryLoader("data/", glob="**/*.pdf", loader_cls=PyPDFLoader)
-    markdown_loader = DirectoryLoader("data/", glob="**/*.md", loader_cls=UnstructuredMarkdownLoader)
-    text_loader = DirectoryLoader("data/", glob="**/*.txt", loader_cls=TextLoader)
-    csv_loader = DirectoryLoader("data/", glob="**/*.csv", loader_cls=CSVLoader)
-    python_loader = DirectoryLoader("data/", glob="**/*.py", loader_cls=PythonLoader)
-    epub_loader = DirectoryLoader("data/", glob="**/*.epub", loader_cls=UnstructuredEPubLoader)
-    html_loader = DirectoryLoader("data/", glob="**/*.html", loader_cls=UnstructuredHTMLLoader)
-    ppt_loader = DirectoryLoader("data/", glob="**/*.ppt", loader_cls=UnstructuredPowerPointLoader)
-    pptx_loader = DirectoryLoader("data/", glob="**/*.pptx", loader_cls=UnstructuredPowerPointLoader)
-    doc_loader = DirectoryLoader("data/", glob="**/*.doc", loader_cls=UnstructuredWordDocumentLoader)
-    docx_loader = DirectoryLoader("data/", glob="**/*.docx", loader_cls=UnstructuredWordDocumentLoader)
-    odt_loader = DirectoryLoader("data/", glob="**/*.odt", loader_cls=UnstructuredODTLoader)
-    notebook_loader = DirectoryLoader("data/", glob="**/*.ipynb", loader_cls=NotebookLoader)
-
-
-    all_loaders = [pdf_loader, markdown_loader, text_loader, csv_loader, python_loader, epub_loader, html_loader, ppt_loader, pptx_loader, doc_loader, docx_loader, odt_loader, notebook_loader]
+    loaders = {
+        "pdf": PyPDFLoader,
+        "md": UnstructuredMarkdownLoader,
+        "txt": TextLoader,
+        "csv": CSVLoader,
+        "py": PythonLoader,
+        "epub": UnstructuredEPubLoader,
+        "html": UnstructuredHTMLLoader,
+        "ppt": UnstructuredPowerPointLoader,
+        "pptx": UnstructuredPowerPointLoader,
+        "doc": UnstructuredWordDocumentLoader,
+        "docx": UnstructuredWordDocumentLoader,
+        "odt": UnstructuredODTLoader,
+        "ipynb": NotebookLoader
+    }
+    # pdf_loader = DirectoryLoader("data/", glob="**/*.pdf", loader_cls=PyPDFLoader)
+    # markdown_loader = DirectoryLoader("data/", glob="**/*.md", loader_cls=UnstructuredMarkdownLoader)
+    # text_loader = DirectoryLoader("data/", glob="**/*.txt", loader_cls=TextLoader)
+    # csv_loader = DirectoryLoader("data/", glob="**/*.csv", loader_cls=CSVLoader)
+    # python_loader = DirectoryLoader("data/", glob="**/*.py", loader_cls=PythonLoader)
+    # epub_loader = DirectoryLoader("data/", glob="**/*.epub", loader_cls=UnstructuredEPubLoader)
+    # html_loader = DirectoryLoader("data/", glob="**/*.html", loader_cls=UnstructuredHTMLLoader)
+    # ppt_loader = DirectoryLoader("data/", glob="**/*.ppt", loader_cls=UnstructuredPowerPointLoader)
+    # pptx_loader = DirectoryLoader("data/", glob="**/*.pptx", loader_cls=UnstructuredPowerPointLoader)
+    # doc_loader = DirectoryLoader("data/", glob="**/*.doc", loader_cls=UnstructuredWordDocumentLoader)
+    # docx_loader = DirectoryLoader("data/", glob="**/*.docx", loader_cls=UnstructuredWordDocumentLoader)
+    # odt_loader = DirectoryLoader("data/", glob="**/*.odt", loader_cls=UnstructuredODTLoader)
+    # notebook_loader = DirectoryLoader("data/", glob="**/*.ipynb", loader_cls=NotebookLoader)
+
+    # Load documents from uploaded files using the appropriate loaders
+    loaded_documents = []
+    for uploaded_file in uploaded_files:
+    # file_extension = os.path.splitext(uploaded_file.name)[-1].lower()[1:]
+    file_extension = os.path.splitext(uploaded_file.name)[-1][1:].lower()
+    if file_extension in loaders:
+        loader_cls = loaders[file_extension]
+        loader = loader_cls()
+        loaded_documents.extend(loader.load(uploaded_file))
+
+    # all_loaders = [pdf_loader, markdown_loader, text_loader, csv_loader, python_loader, epub_loader, html_loader, ppt_loader, pptx_loader, doc_loader, docx_loader, odt_loader, notebook_loader]
 
     # Load documents from all loaders
-    loaded_documents = []
-    for loader in all_loaders:
-        loaded_documents.extend(loader.load())
+    # for loader in all_loaders:
+    #     loaded_documents.extend(loader.load())
 
     # Split loaded documents into chunks
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=40)