init commit

.chainlit/translations/en-US.json

@@ -1,231 +0,0 @@
-{
-    "components": {
-        "atoms": {
-            "buttons": {
-                "userButton": {
-                    "menu": {
-                        "settings": "Settings",
-                        "settingsKey": "S",
-                        "APIKeys": "API Keys",
-                        "logout": "Logout"
-                    }
-                }
-            }
-        },
-        "molecules": {
-            "newChatButton": {
-                "newChat": "New Chat"
-            },
-            "tasklist": {
-                "TaskList": {
-                    "title": "\ud83d\uddd2\ufe0f Task List",
-                    "loading": "Loading...",
-                    "error": "An error occured"
-                }
-            },
-            "attachments": {
-                "cancelUpload": "Cancel upload",
-                "removeAttachment": "Remove attachment"
-            },
-            "newChatDialog": {
-                "createNewChat": "Create new chat?",
-                "clearChat": "This will clear the current messages and start a new chat.",
-                "cancel": "Cancel",
-                "confirm": "Confirm"
-            },
-            "settingsModal": {
-                "settings": "Settings",
-                "expandMessages": "Expand Messages",
-                "hideChainOfThought": "Hide Chain of Thought",
-                "darkMode": "Dark Mode"
-            },
-            "detailsButton": {
-                "using": "Using",
-                "running": "Running",
-                "took_one": "Took {{count}} step",
-                "took_other": "Took {{count}} steps"
-            },
-            "auth": {
-                "authLogin": {
-                    "title": "Login to access the app.",
-                    "form": {
-                        "email": "Email address",
-                        "password": "Password",
-                        "noAccount": "Don't have an account?",
-                        "alreadyHaveAccount": "Already have an account?",
-                        "signup": "Sign Up",
-                        "signin": "Sign In",
-                        "or": "OR",
-                        "continue": "Continue",
-                        "forgotPassword": "Forgot password?",
-                        "passwordMustContain": "Your password must contain:",
-                        "emailRequired": "email is a required field",
-                        "passwordRequired": "password is a required field"
-                    },
-                    "error": {
-                        "default": "Unable to sign in.",
-                        "signin": "Try signing in with a different account.",
-                        "oauthsignin": "Try signing in with a different account.",
-                        "redirect_uri_mismatch": "The redirect URI is not matching the oauth app configuration.",
-                        "oauthcallbackerror": "Try signing in with a different account.",
-                        "oauthcreateaccount": "Try signing in with a different account.",
-                        "emailcreateaccount": "Try signing in with a different account.",
-                        "callback": "Try signing in with a different account.",
-                        "oauthaccountnotlinked": "To confirm your identity, sign in with the same account you used originally.",
-                        "emailsignin": "The e-mail could not be sent.",
-                        "emailverify": "Please verify your email, a new email has been sent.",
-                        "credentialssignin": "Sign in failed. Check the details you provided are correct.",
-                        "sessionrequired": "Please sign in to access this page."
-                    }
-                },
-                "authVerifyEmail": {
-                    "almostThere": "You're almost there! We've sent an email to ",
-                    "verifyEmailLink": "Please click on the link in that email to complete your signup.",
-                    "didNotReceive": "Can't find the email?",
-                    "resendEmail": "Resend email",
-                    "goBack": "Go Back",
-                    "emailSent": "Email sent successfully.",
-                    "verifyEmail": "Verify your email address"
-                },
-                "providerButton": {
-                    "continue": "Continue with {{provider}}",
-                    "signup": "Sign up with {{provider}}"
-                },
-                "authResetPassword": {
-                    "newPasswordRequired": "New password is a required field",
-                    "passwordsMustMatch": "Passwords must match",
-                    "confirmPasswordRequired": "Confirm password is a required field",
-                    "newPassword": "New password",
-                    "confirmPassword": "Confirm password",
-                    "resetPassword": "Reset Password"
-                },
-                "authForgotPassword": {
-                    "email": "Email address",
-                    "emailRequired": "email is a required field",
-                    "emailSent": "Please check the email address {{email}} for instructions to reset your password.",
-                    "enterEmail": "Enter your email address and we will send you instructions to reset your password.",
-                    "resendEmail": "Resend email",
-                    "continue": "Continue",
-                    "goBack": "Go Back"
-                }
-            }
-        },
-        "organisms": {
-            "chat": {
-                "history": {
-                    "index": {
-                        "showHistory": "Show history",
-                        "lastInputs": "Last Inputs",
-                        "noInputs": "Such empty...",
-                        "loading": "Loading..."
-                    }
-                },
-                "inputBox": {
-                    "input": {
-                        "placeholder": "Type your message here..."
-                    },
-                    "speechButton": {
-                        "start": "Start recording",
-                        "stop": "Stop recording"
-                    },
-                    "SubmitButton": {
-                        "sendMessage": "Send message",
-                        "stopTask": "Stop Task"
-                    },
-                    "UploadButton": {
-                        "attachFiles": "Attach files"
-                    },
-                    "waterMark": {
-                        "text": "Built with"
-                    }
-                },
-                "Messages": {
-                    "index": {
-                        "running": "Running",
-                        "executedSuccessfully": "executed successfully",
-                        "failed": "failed",
-                        "feedbackUpdated": "Feedback updated",
-                        "updating": "Updating"
-                    }
-                },
-                "dropScreen": {
-                    "dropYourFilesHere": "Drop your files here"
-                },
-                "index": {
-                    "failedToUpload": "Failed to upload",
-                    "cancelledUploadOf": "Cancelled upload of",
-                    "couldNotReachServer": "Could not reach the server",
-                    "continuingChat": "Continuing previous chat"
-                },
-                "settings": {
-                    "settingsPanel": "Settings panel",
-                    "reset": "Reset",
-                    "cancel": "Cancel",
-                    "confirm": "Confirm"
-                }
-            },
-            "threadHistory": {
-                "sidebar": {
-                    "filters": {
-                        "FeedbackSelect": {
-                            "feedbackAll": "Feedback: All",
-                            "feedbackPositive": "Feedback: Positive",
-                            "feedbackNegative": "Feedback: Negative"
-                        },
-                        "SearchBar": {
-                            "search": "Search"
-                        }
-                    },
-                    "DeleteThreadButton": {
-                        "confirmMessage": "This will delete the thread as well as it's messages and elements.",
-                        "cancel": "Cancel",
-                        "confirm": "Confirm",
-                        "deletingChat": "Deleting chat",
-                        "chatDeleted": "Chat deleted"
-                    },
-                    "index": {
-                        "pastChats": "Past Chats"
-                    },
-                    "ThreadList": {
-                        "empty": "Empty...",
-                        "today": "Today",
-                        "yesterday": "Yesterday",
-                        "previous7days": "Previous 7 days",
-                        "previous30days": "Previous 30 days"
-                    },
-                    "TriggerButton": {
-                        "closeSidebar": "Close sidebar",
-                        "openSidebar": "Open sidebar"
-                    }
-                },
-                "Thread": {
-                    "backToChat": "Go back to chat",
-                    "chatCreatedOn": "This chat was created on"
-                }
-            },
-            "header": {
-                "chat": "Chat",
-                "readme": "Readme"
-            }
-        }
-    },
-    "hooks": {
-        "useLLMProviders": {
-            "failedToFetchProviders": "Failed to fetch providers:"
-        }
-    },
-    "pages": {
-        "Design": {},
-        "Env": {
-            "savedSuccessfully": "Saved successfully",
-            "requiredApiKeys": "Required API Keys",
-            "requiredApiKeysInfo": "To use this app, the following API keys are required. The keys are stored on your device's local storage."
-        },
-        "Page": {
-            "notPartOfProject": "You are not part of this project."
-        },
-        "ResumeButton": {
-            "resumeChat": "Resume Chat"
-        }
-    }
-}

.chainlit/translations/pt-BR.json

@@ -1,155 +0,0 @@
-{
-    "components": {
-        "atoms": {
-            "buttons": {
-                "userButton": {
-                    "menu": {
-                        "settings": "Configura\u00e7\u00f5es",
-                        "settingsKey": "S",
-                        "APIKeys": "Chaves de API",
-                        "logout": "Sair"
-                    }
-                }
-            }
-        },
-        "molecules": {
-            "newChatButton": {
-                "newChat": "Nova Conversa"
-            },
-            "tasklist": {
-                "TaskList": {
-                    "title": "\ud83d\uddd2\ufe0f Lista de Tarefas",
-                    "loading": "Carregando...",
-                    "error": "Ocorreu um erro"
-                }
-            },
-            "attachments": {
-                "cancelUpload": "Cancelar envio",
-                "removeAttachment": "Remover anexo"
-            },
-            "newChatDialog": {
-                "createNewChat": "Criar novo chat?",
-                "clearChat": "Isso limpar\u00e1 as mensagens atuais e iniciar\u00e1 uma nova conversa.",
-                "cancel": "Cancelar",
-                "confirm": "Confirmar"
-            },
-            "settingsModal": {
-                "expandMessages": "Expandir Mensagens",
-                "hideChainOfThought": "Esconder Sequ\u00eancia de Pensamento",
-                "darkMode": "Modo Escuro"
-            }
-        },
-        "organisms": {
-            "chat": {
-                "history": {
-                    "index": {
-                        "lastInputs": "\u00daltimas Entradas",
-                        "noInputs": "Vazio...",
-                        "loading": "Carregando..."
-                    }
-                },
-                "inputBox": {
-                    "input": {
-                        "placeholder": "Digite sua mensagem aqui..."
-                    },
-                    "speechButton": {
-                        "start": "Iniciar grava\u00e7\u00e3o",
-                        "stop": "Parar grava\u00e7\u00e3o"
-                    },
-                    "SubmitButton": {
-                        "sendMessage": "Enviar mensagem",
-                        "stopTask": "Parar Tarefa"
-                    },
-                    "UploadButton": {
-                        "attachFiles": "Anexar arquivos"
-                    },
-                    "waterMark": {
-                        "text": "Constru\u00eddo com"
-                    }
-                },
-                "Messages": {
-                    "index": {
-                        "running": "Executando",
-                        "executedSuccessfully": "executado com sucesso",
-                        "failed": "falhou",
-                        "feedbackUpdated": "Feedback atualizado",
-                        "updating": "Atualizando"
-                    }
-                },
-                "dropScreen": {
-                    "dropYourFilesHere": "Solte seus arquivos aqui"
-                },
-                "index": {
-                    "failedToUpload": "Falha ao enviar",
-                    "cancelledUploadOf": "Envio cancelado de",
-                    "couldNotReachServer": "N\u00e3o foi poss\u00edvel conectar ao servidor",
-                    "continuingChat": "Continuando o chat anterior"
-                },
-                "settings": {
-                    "settingsPanel": "Painel de Configura\u00e7\u00f5es",
-                    "reset": "Redefinir",
-                    "cancel": "Cancelar",
-                    "confirm": "Confirmar"
-                }
-            },
-            "threadHistory": {
-                "sidebar": {
-                    "filters": {
-                        "FeedbackSelect": {
-                            "feedbackAll": "Feedback: Todos",
-                            "feedbackPositive": "Feedback: Positivo",
-                            "feedbackNegative": "Feedback: Negativo"
-                        },
-                        "SearchBar": {
-                            "search": "Buscar"
-                        }
-                    },
-                    "DeleteThreadButton": {
-                        "confirmMessage": "Isso deletar\u00e1 a conversa, assim como suas mensagens e elementos.",
-                        "cancel": "Cancelar",
-                        "confirm": "Confirmar",
-                        "deletingChat": "Deletando conversa",
-                        "chatDeleted": "Conversa deletada"
-                    },
-                    "index": {
-                        "pastChats": "Conversas Anteriores"
-                    },
-                    "ThreadList": {
-                        "empty": "Vazio..."
-                    },
-                    "TriggerButton": {
-                        "closeSidebar": "Fechar barra lateral",
-                        "openSidebar": "Abrir barra lateral"
-                    }
-                },
-                "Thread": {
-                    "backToChat": "Voltar para a conversa",
-                    "chatCreatedOn": "Esta conversa foi criada em"
-                }
-            },
-            "header": {
-                "chat": "Conversa",
-                "readme": "Leia-me"
-            }
-        },
-        "hooks": {
-            "useLLMProviders": {
-                "failedToFetchProviders": "Falha ao buscar provedores:"
-            }
-        },
-        "pages": {
-            "Design": {},
-            "Env": {
-                "savedSuccessfully": "Salvo com sucesso",
-                "requiredApiKeys": "Chaves de API necess\u00e1rias",
-                "requiredApiKeysInfo": "Para usar este aplicativo, as seguintes chaves de API s\u00e3o necess\u00e1rias. As chaves s\u00e3o armazenadas localmente em seu dispositivo."
-            },
-            "Page": {
-                "notPartOfProject": "Voc\u00ea n\u00e3o faz parte deste projeto."
-            },
-            "ResumeButton": {
-                "resumeChat": "Continuar Conversa"
-            }
-        }
-    }
-}

.gitignore

@@ -160,4 +160,10 @@ cython_debug/
 #.idea/
 
 # log files
-*.log
+*.log
+
+.ragatouille/*
+*/__pycache__/*
+*/.chainlit/translations/*
+storage/logs/*
+vectorstores/*

code/__init__.py

@@ -0,0 +1 @@
+from .modules import *

code/main.py

@@ -1,9 +1,8 @@
-from langchain.document_loaders import PyPDFLoader, DirectoryLoader
+from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
 from langchain import PromptTemplate
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
 from langchain.chains import RetrievalQA
-from langchain.llms import CTransformers
 import chainlit as cl
 from langchain_community.chat_models import ChatOpenAI
 from langchain_community.embeddings import OpenAIEmbeddings
@@ -11,13 +10,22 @@ import yaml
 import logging
 from dotenv import load_dotenv
 
-from modules.llm_tutor import LLMTutor
-from modules.constants import *
-from modules.helpers import get_sources
+import os
+import sys
 
+# Add the 'code' directory to the Python path
+current_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(current_dir)
 
+from modules.chat.llm_tutor import LLMTutor
+from modules.config.constants import *
+from modules.chat.helpers import get_sources
+
+
+global logger
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
+logger.propagate = False
 
 # Console Handler
 console_handler = logging.StreamHandler()
@@ -26,13 +34,6 @@ formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
 console_handler.setFormatter(formatter)
 logger.addHandler(console_handler)
 
-# File Handler
-log_file_path = "log_file.log"  # Change this to your desired log file path
-file_handler = logging.FileHandler(log_file_path)
-file_handler.setLevel(logging.INFO)
-file_handler.setFormatter(formatter)
-logger.addHandler(file_handler)
-
 
 # Adding option to select the chat profile
 @cl.set_chat_profiles
@@ -66,12 +67,26 @@ def rename(orig_author: str):
 # chainlit code
 @cl.on_chat_start
 async def start():
-    with open("code/config.yml", "r") as f:
+    with open("modules/config/config.yml", "r") as f:
         config = yaml.safe_load(f)
-        print(config)
-        logger.info("Config file loaded")
-        logger.info(f"Config: {config}")
-        logger.info("Creating llm_tutor instance")
+
+    # Ensure log directory exists
+    log_directory = config["log_dir"]
+    if not os.path.exists(log_directory):
+        os.makedirs(log_directory)
+
+    # File Handler
+    log_file_path = (
+        f"{log_directory}/tutor.log"  # Change this to your desired log file path
+    )
+    file_handler = logging.FileHandler(log_file_path, mode="w")
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+
+    logger.info("Config file loaded")
+    logger.info(f"Config: {config}")
+    logger.info("Creating llm_tutor instance")
 
     chat_profile = cl.user_session.get("chat_profile")
     if chat_profile is not None:
@@ -93,8 +108,7 @@ async def start():
     llm_tutor = LLMTutor(config, logger=logger)
 
     chain = llm_tutor.qa_bot()
-    model = config["llm_params"]["local_llm_params"]["model"]
-    msg = cl.Message(content=f"Starting the bot {model}...")
+    msg = cl.Message(content=f"Starting the bot {chat_profile}...")
     await msg.send()
     msg.content = opening_message
     await msg.update()
@@ -104,24 +118,17 @@ async def start():
 
 @cl.on_message
 async def main(message):
+    global logger
     user = cl.user_session.get("user")
     chain = cl.user_session.get("chain")
-    # cb = cl.AsyncLangchainCallbackHandler(
-    #     stream_final_answer=True, answer_prefix_tokens=["FINAL", "ANSWER"]
-    # )
-    # cb.answer_reached = True
-    # res=await chain.acall(message, callbacks=[cb])
-    res = await chain.acall(message.content)
-    print(f"response: {res}")
+    cb = cl.AsyncLangchainCallbackHandler()  # TODO: fix streaming here
+    cb.answer_reached = True
+    res = await chain.acall(message.content, callbacks=[cb])
+    # res = await chain.acall(message.content)
     try:
         answer = res["answer"]
     except:
         answer = res["result"]
-    print(f"answer: {answer}")
-
-    logger.info(f"Question: {res['question']}")
-    logger.info(f"History: {res['chat_history']}")
-    logger.info(f"Answer: {answer}\n")
 
     answer_with_sources, source_elements = get_sources(res, answer)
 

code/modules/__init__.py

@@ -0,0 +1,2 @@
+from . import vectorstore
+from . import dataloader

code/modules/chat/__init__.py

@@ -0,0 +1,2 @@
+from .llm_tutor import LLMTutor
+from .chat_model_loader import ChatModelLoader

code/modules/{chat_model_loader.py → chat/chat_model_loader.py}

@@ -1,8 +1,7 @@
 from langchain_community.chat_models import ChatOpenAI
-from langchain.llms import CTransformers
-from langchain.llms.huggingface_pipeline import HuggingFacePipeline
+from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
 from transformers import AutoTokenizer, TextStreamer
-from langchain.llms import LlamaCpp
+from langchain_community.llms import LlamaCpp
 import torch
 import transformers
 import os

code/modules/{helpers.py → chat/helpers.py}

@@ -1,176 +1,6 @@
-import requests
-from bs4 import BeautifulSoup
-from tqdm import tqdm
+from modules.config.constants import *
 import chainlit as cl
-from langchain import PromptTemplate
-import requests
-from bs4 import BeautifulSoup
-from urllib.parse import urlparse, urljoin, urldefrag
-import asyncio
-import aiohttp
-from aiohttp import ClientSession
-from typing import Dict, Any, List
-
-try:
-    from modules.constants import *
-except:
-    from constants import *
-
-"""
-Ref: https://python.plainenglish.io/scraping-the-subpages-on-a-website-ea2d4e3db113
-"""
-
-
-class WebpageCrawler:
-    def __init__(self):
-        self.dict_href_links = {}
-
-    async def fetch(self, session: ClientSession, url: str) -> str:
-        async with session.get(url) as response:
-            try:
-                return await response.text()
-            except UnicodeDecodeError:
-                return await response.text(encoding="latin1")
-
-    def url_exists(self, url: str) -> bool:
-        try:
-            response = requests.head(url)
-            return response.status_code == 200
-        except requests.ConnectionError:
-            return False
-
-    async def get_links(self, session: ClientSession, website_link: str, base_url: str):
-        html_data = await self.fetch(session, website_link)
-        soup = BeautifulSoup(html_data, "html.parser")
-        list_links = []
-        for link in soup.find_all("a", href=True):
-            href = link["href"].strip()
-            full_url = urljoin(base_url, href)
-            normalized_url = self.normalize_url(full_url)  # sections removed
-            if (
-                normalized_url not in self.dict_href_links
-                and self.is_child_url(normalized_url, base_url)
-                and self.url_exists(normalized_url)
-            ):
-                self.dict_href_links[normalized_url] = None
-                list_links.append(normalized_url)
-
-        return list_links
-
-    async def get_subpage_links(
-        self, session: ClientSession, urls: list, base_url: str
-    ):
-        tasks = [self.get_links(session, url, base_url) for url in urls]
-        results = await asyncio.gather(*tasks)
-        all_links = [link for sublist in results for link in sublist]
-        return all_links
-
-    async def get_all_pages(self, url: str, base_url: str):
-        async with aiohttp.ClientSession() as session:
-            dict_links = {url: "Not-checked"}
-            counter = None
-            while counter != 0:
-                unchecked_links = [
-                    link
-                    for link, status in dict_links.items()
-                    if status == "Not-checked"
-                ]
-                if not unchecked_links:
-                    break
-                new_links = await self.get_subpage_links(
-                    session, unchecked_links, base_url
-                )
-                for link in unchecked_links:
-                    dict_links[link] = "Checked"
-                    print(f"Checked: {link}")
-                dict_links.update(
-                    {
-                        link: "Not-checked"
-                        for link in new_links
-                        if link not in dict_links
-                    }
-                )
-                counter = len(
-                    [
-                        status
-                        for status in dict_links.values()
-                        if status == "Not-checked"
-                    ]
-                )
-
-            checked_urls = [
-                url for url, status in dict_links.items() if status == "Checked"
-            ]
-            return checked_urls
-
-    def is_webpage(self, url: str) -> bool:
-        try:
-            response = requests.head(url, allow_redirects=True)
-            content_type = response.headers.get("Content-Type", "").lower()
-            return "text/html" in content_type
-        except requests.RequestException:
-            return False
-
-    def clean_url_list(self, urls):
-        files, webpages = [], []
-
-        for url in urls:
-            if self.is_webpage(url):
-                webpages.append(url)
-            else:
-                files.append(url)
-
-        return files, webpages
-
-    def is_child_url(self, url, base_url):
-        return url.startswith(base_url)
-
-    def normalize_url(self, url: str):
-        # Strip the fragment identifier
-        defragged_url, _ = urldefrag(url)
-        return defragged_url
-
-
-def get_urls_from_file(file_path: str):
-    """
-    Function to get urls from a file
-    """
-    with open(file_path, "r") as f:
-        urls = f.readlines()
-    urls = [url.strip() for url in urls]
-    return urls
-
-
-def get_base_url(url):
-    parsed_url = urlparse(url)
-    base_url = f"{parsed_url.scheme}://{parsed_url.netloc}/"
-    return base_url
-
-
-def get_prompt(config):
-    if config["llm_params"]["use_history"]:
-        if config["llm_params"]["llm_loader"] == "local_llm":
-            custom_prompt_template = tinyllama_prompt_template_with_history
-        elif config["llm_params"]["llm_loader"] == "openai":
-            custom_prompt_template = openai_prompt_template_with_history
-        # else:
-        #     custom_prompt_template = tinyllama_prompt_template_with_history # default
-        prompt = PromptTemplate(
-            template=custom_prompt_template,
-            input_variables=["context", "chat_history", "question"],
-        )
-    else:
-        if config["llm_params"]["llm_loader"] == "local_llm":
-            custom_prompt_template = tinyllama_prompt_template
-        elif config["llm_params"]["llm_loader"] == "openai":
-            custom_prompt_template = openai_prompt_template
-        # else:
-        #     custom_prompt_template = tinyllama_prompt_template
-        prompt = PromptTemplate(
-            template=custom_prompt_template,
-            input_variables=["context", "question"],
-        )
-    return prompt
+from langchain_core.prompts import PromptTemplate
 
 
 def get_sources(res, answer):
@@ -248,90 +78,27 @@ def get_sources(res, answer):
     return full_answer, source_elements
 
 
-def get_metadata(lectures_url, schedule_url):
-    """
-    Function to get the lecture metadata from the lectures and schedule URLs.
-    """
-    lecture_metadata = {}
-
-    # Get the main lectures page content
-    r_lectures = requests.get(lectures_url)
-    soup_lectures = BeautifulSoup(r_lectures.text, "html.parser")
-
-    # Get the main schedule page content
-    r_schedule = requests.get(schedule_url)
-    soup_schedule = BeautifulSoup(r_schedule.text, "html.parser")
-
-    # Find all lecture blocks
-    lecture_blocks = soup_lectures.find_all("div", class_="lecture-container")
-
-    # Create a mapping from slides link to date
-    date_mapping = {}
-    schedule_rows = soup_schedule.find_all("li", class_="table-row-lecture")
-    for row in schedule_rows:
-        try:
-            date = (
-                row.find("div", {"data-label": "Date"}).get_text(separator=" ").strip()
-            )
-            description_div = row.find("div", {"data-label": "Description"})
-            slides_link_tag = description_div.find("a", title="Download slides")
-            slides_link = slides_link_tag["href"].strip() if slides_link_tag else None
-            slides_link = (
-                f"https://dl4ds.github.io{slides_link}" if slides_link else None
-            )
-            if slides_link:
-                date_mapping[slides_link] = date
-        except Exception as e:
-            print(f"Error processing schedule row: {e}")
-            continue
-
-    for block in lecture_blocks:
-        try:
-            # Extract the lecture title
-            title = block.find("span", style="font-weight: bold;").text.strip()
-
-            # Extract the TL;DR
-            tldr = block.find("strong", text="tl;dr:").next_sibling.strip()
-
-            # Extract the link to the slides
-            slides_link_tag = block.find("a", title="Download slides")
-            slides_link = slides_link_tag["href"].strip() if slides_link_tag else None
-            slides_link = (
-                f"https://dl4ds.github.io{slides_link}" if slides_link else None
-            )
-
-            # Extract the link to the lecture recording
-            recording_link_tag = block.find("a", title="Download lecture recording")
-            recording_link = (
-                recording_link_tag["href"].strip() if recording_link_tag else None
-            )
-
-            # Extract suggested readings or summary if available
-            suggested_readings_tag = block.find("p", text="Suggested Readings:")
-            if suggested_readings_tag:
-                suggested_readings = suggested_readings_tag.find_next_sibling("ul")
-                if suggested_readings:
-                    suggested_readings = suggested_readings.get_text(
-                        separator="\n"
-                    ).strip()
-                else:
-                    suggested_readings = "No specific readings provided."
-            else:
-                suggested_readings = "No specific readings provided."
-
-            # Get the date from the schedule
-            date = date_mapping.get(slides_link, "No date available")
-
-            # Add to the dictionary
-            lecture_metadata[slides_link] = {
-                "date": date,
-                "tldr": tldr,
-                "title": title,
-                "lecture_recording": recording_link,
-                "suggested_readings": suggested_readings,
-            }
-        except Exception as e:
-            print(f"Error processing block: {e}")
-            continue
-
-    return lecture_metadata
+def get_prompt(config):
+    if config["llm_params"]["use_history"]:
+        if config["llm_params"]["llm_loader"] == "local_llm":
+            custom_prompt_template = tinyllama_prompt_template_with_history
+        elif config["llm_params"]["llm_loader"] == "openai":
+            custom_prompt_template = openai_prompt_template_with_history
+        # else:
+        #     custom_prompt_template = tinyllama_prompt_template_with_history # default
+        prompt = PromptTemplate(
+            template=custom_prompt_template,
+            input_variables=["context", "chat_history", "question"],
+        )
+    else:
+        if config["llm_params"]["llm_loader"] == "local_llm":
+            custom_prompt_template = tinyllama_prompt_template
+        elif config["llm_params"]["llm_loader"] == "openai":
+            custom_prompt_template = openai_prompt_template
+        # else:
+        #     custom_prompt_template = tinyllama_prompt_template
+        prompt = PromptTemplate(
+            template=custom_prompt_template,
+            input_variables=["context", "question"],
+        )
+    return prompt

code/modules/{llm_tutor.py → chat/llm_tutor.py}

@@ -1,24 +1,52 @@
-from langchain import PromptTemplate
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain_community.chat_models import ChatOpenAI
-from langchain_community.embeddings import OpenAIEmbeddings
-from langchain.vectorstores import FAISS
 from langchain.chains import RetrievalQA, ConversationalRetrievalChain
-from langchain.llms import CTransformers
-from langchain.memory import ConversationBufferWindowMemory, ConversationSummaryBufferMemory
+from langchain.memory import (
+    ConversationBufferWindowMemory,
+    ConversationSummaryBufferMemory,
+)
 from langchain.chains.conversational_retrieval.prompts import QA_PROMPT
 import os
-from modules.constants import *
-from modules.helpers import get_prompt
-from modules.chat_model_loader import ChatModelLoader
-from modules.vector_db import VectorDB, VectorDBScore
-from typing import Dict, Any, Optional
+from modules.config.constants import *
+from modules.chat.helpers import get_prompt
+from modules.chat.chat_model_loader import ChatModelLoader
+from modules.vectorstore.store_manager import VectorStoreManager
+
+from modules.retriever import FaissRetriever, ChromaRetriever
+
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 from langchain_core.callbacks.manager import AsyncCallbackManagerForChainRun
 import inspect
 from langchain.chains.conversational_retrieval.base import _get_chat_history
+from langchain_core.messages import BaseMessage
+
+CHAT_TURN_TYPE = Union[Tuple[str, str], BaseMessage]
+
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_community.chat_models import ChatOpenAI
 
 
 class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
+
+    def _get_chat_history(self, chat_history: List[CHAT_TURN_TYPE]) -> str:
+        _ROLE_MAP = {"human": "Student: ", "ai": "AI Tutor: "}
+        buffer = ""
+        for dialogue_turn in chat_history:
+            if isinstance(dialogue_turn, BaseMessage):
+                role_prefix = _ROLE_MAP.get(
+                    dialogue_turn.type, f"{dialogue_turn.type}: "
+                )
+                buffer += f"\n{role_prefix}{dialogue_turn.content}"
+            elif isinstance(dialogue_turn, tuple):
+                human = "Student: " + dialogue_turn[0]
+                ai = "AI Tutor: " + dialogue_turn[1]
+                buffer += "\n" + "\n".join([human, ai])
+            else:
+                raise ValueError(
+                    f"Unsupported chat history format: {type(dialogue_turn)}."
+                    f" Full chat history: {chat_history} "
+                )
+        return buffer
+
     async def _acall(
         self,
         inputs: Dict[str, Any],
@@ -26,13 +54,31 @@ class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
     ) -> Dict[str, Any]:
         _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()
         question = inputs["question"]
-        get_chat_history = self.get_chat_history or _get_chat_history
+        get_chat_history = self._get_chat_history
         chat_history_str = get_chat_history(inputs["chat_history"])
-        print(f"chat_history_str: {chat_history_str}")
         if chat_history_str:
-            callbacks = _run_manager.get_child()
-            new_question = await self.question_generator.arun(
-                question=question, chat_history=chat_history_str, callbacks=callbacks
+            # callbacks = _run_manager.get_child()
+            # new_question = await self.question_generator.arun(
+            #     question=question, chat_history=chat_history_str, callbacks=callbacks
+            # )
+            system = (
+                "You are an AI Tutor helping a student. Your task is to rephrase the student's question to provide more context from their chat history (only if relevant), ensuring the rephrased question still reflects the student's point of view. "
+                "The rephrased question should incorporate relevant details from the chat history to make it clearer and more specific. It should also expand upon the original question to provide more context on only what the student provided."
+                "Always end the rephrased question with the original question in parentheses for reference. "
+                "Do not change the meaning of the question, and keep the tone and perspective as if it were asked by the student. "
+                "Here is the chat history for context: \n{chat_history_str}\n"
+                "Now, rephrase the following question: '{question}'"
+            )
+            prompt = ChatPromptTemplate.from_messages(
+                [
+                    ("system", system),
+                    ("human", "{question}, {chat_history_str}"),
+                ]
+            )
+            llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
+            step_back = prompt | llm | StrOutputParser()
+            new_question = step_back.invoke(
+                {"question": question, "chat_history_str": chat_history_str}
             )
         else:
             new_question = question
@@ -56,27 +102,24 @@ class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
             # Prepare the final prompt with metadata
             context = "\n\n".join(
                 [
-                    f"Document content: {doc.page_content}\nMetadata: {doc.metadata}"
-                    for doc in docs
+                    f"Context {idx+1}: \n(Document content: {doc.page_content}\nMetadata: (source_file: {doc.metadata['source']}))"
+                    for idx, doc in enumerate(docs)
                 ]
             )
-            final_prompt = f"""
-                You are an AI Tutor for the course DS598, taught by Prof. Thomas Gardos. Use the following pieces of information to answer the user's question. 
-                If you don't know the answer, just say that you don't know—don't try to make up an answer. 
-                Use the chat history to answer the question only if it's relevant; otherwise, ignore it. The context for the answer will be under "Document context:". 
-                Use the metadata from each document to guide the user to the correct sources. 
-                The context is ordered by relevance to the question. Give more weight to the most relevant documents.
-                Talk in a friendly and personalized manner, similar to how you would speak to a friend who needs help. Make the conversation engaging and avoid sounding repetitive or robotic.
-
-                Chat History:
-                {chat_history_str}
-
-                Context:
-                {context}
-
-                Question: {new_question}
-                AI Tutor:
-                """
+            final_prompt = (
+                "You are an AI Tutor for the course DS598, taught by Prof. Thomas Gardos. "
+                "Use the following pieces of information to answer the user's question. "
+                "If you don't know the answer, try your best, but don't try to make up an answer. Keep the flow of the conversation going. "
+                "Use the chat history just as a gist to answer the question only if it's relevant; otherwise, ignore it. Do not repeat responses in the history. Use the context as a guide to construct your answer. The context for the answer will be under 'Document context:'. Remember, the conext may include text not directly related to the question."
+                "Make sure to use the source_file field in metadata from each document to provide links to the user to the correct sources. "
+                "The context is ordered by relevance to the question. "
+                "Talk in a friendly and personalized manner, similar to how you would speak to a friend who needs help. Make the conversation engaging and avoid sounding repetitive or robotic.\n\n"
+                f"Chat History:\n{chat_history_str}\n\n"
+                f"Context:\n{context}\n\n"
+                f"Student: {new_question}\n"
+                "Anwer the student's question in a friendly, concise, and engaging manner.\n"
+                "AI Tutor:"
+            )
 
             new_inputs["input"] = final_prompt
             new_inputs["question"] = final_prompt
@@ -98,8 +141,9 @@ class LLMTutor:
     def __init__(self, config, logger=None):
         self.config = config
         self.llm = self.load_llm()
-        self.vector_db = VectorDB(config, logger=logger)
-        if self.config["embedding_options"]["embedd_files"]:
+        self.logger = logger
+        self.vector_db = VectorStoreManager(config, logger=self.logger)
+        if self.config["vectorstore"]["embedd_files"]:
             self.vector_db.create_database()
             self.vector_db.save_database()
 
@@ -114,24 +158,20 @@ class LLMTutor:
 
     # Retrieval QA Chain
     def retrieval_qa_chain(self, llm, prompt, db):
-        if self.config["embedding_options"]["db_option"] in ["FAISS", "Chroma"]:
-            retriever = VectorDBScore(
-                vectorstore=db,
-                # search_type="similarity_score_threshold",
-                # search_kwargs={
-                #     "score_threshold": self.config["embedding_options"][
-                #         "score_threshold"
-                #     ],
-                #     "k": self.config["embedding_options"]["search_top_k"],
-                # },
-            )
-        elif self.config["embedding_options"]["db_option"] == "RAGatouille":
+
+        if self.config["vectorstore"]["db_option"] == "FAISS":
+            retriever = FaissRetriever().return_retriever(db, self.config)
+
+        elif self.config["vectorstore"]["db_option"] == "Chroma":
+            retriever = ChromaRetriever().return_retriever(db, self.config)
+
+        elif self.config["vectorstore"]["db_option"] == "RAGatouille":
             retriever = db.as_langchain_retriever(
-                k=self.config["embedding_options"]["search_top_k"]
+                k=self.config["vectorstore"]["search_top_k"]
             )
+
         if self.config["llm_params"]["use_history"]:
-            memory = ConversationSummaryBufferMemory(
-                llm = llm,
+            memory = ConversationBufferWindowMemory(
                 k=self.config["llm_params"]["memory_window"],
                 memory_key="chat_history",
                 return_messages=True,
@@ -145,6 +185,7 @@ class LLMTutor:
                 return_source_documents=True,
                 memory=memory,
                 combine_docs_chain_kwargs={"prompt": prompt},
+                response_if_no_docs_found="No context found",
             )
         else:
             qa_chain = RetrievalQA.from_chain_type(
@@ -166,7 +207,9 @@ class LLMTutor:
     def qa_bot(self):
         db = self.vector_db.load_database()
         qa_prompt = self.set_custom_prompt()
-        qa = self.retrieval_qa_chain(self.llm, qa_prompt, db)
+        qa = self.retrieval_qa_chain(
+            self.llm, qa_prompt, db
+        )  # TODO:  PROMPT is overwritten in CustomConversationalRetrievalChain
 
         return qa
 

code/modules/config/__init__.py

@@ -0,0 +1 @@
+from .constants import *

code/{config.yml → modules/config/config.yml}

@@ -1,23 +1,32 @@
-embedding_options:
+log_dir: '../storage/logs' # str
+log_chunk_dir: '../storage/logs/chunks' # str
+device: 'cpu' # str [cuda, cpu]
+
+vectorstore:
   embedd_files: False # bool
-  data_path: 'storage/data' # str
-  url_file_path: 'storage/data/urls.txt' # str
-  expand_urls: True # bool
-  db_option : 'RAGatouille' # str [FAISS, Chroma, RAGatouille]
-  db_path : 'vectorstores' # str
+  data_path: '../storage/data' # str
+  url_file_path: '../storage/data/urls.txt' # str
+  expand_urls: False # bool
+  db_option : 'Chroma' # str [FAISS, Chroma, RAGatouille]
+  db_path : '../vectorstores' # str
   model : 'sentence-transformers/all-MiniLM-L6-v2' # str [sentence-transformers/all-MiniLM-L6-v2, text-embedding-ada-002']
   search_top_k : 3 # int
   score_threshold : 0.2 # float
+
+  faiss_params: # Not used as of now
+    index_path: '../vectorstores/faiss.index' # str
+    index_type: 'Flat' # str [Flat, HNSW, IVF]
+    index_dimension: 384 # int
+    index_nlist: 100 # int
+    index_nprobe: 10 # int
+
 llm_params: 
   use_history: True # bool
   memory_window: 3 # int
   llm_loader: 'openai' # str [local_llm, openai]
   openai_params:
     model: 'gpt-3.5-turbo-1106' # str [gpt-3.5-turbo-1106, gpt-4]
-  local_llm_params:
-    model: "storage/models/llama-2-7b-chat.Q4_0.gguf"
-    model_type: "llama"
-    temperature: 0.2
+
 splitter_options:
   use_splitter: True # bool
   split_by_token : True # bool

code/modules/dataloader/__init__.py

@@ -0,0 +1,2 @@
+from .webpage_crawler import WebpageCrawler
+from .data_loader import DataLoader

code/modules/{data_loader.py → dataloader/data_loader.py}

@@ -16,15 +16,12 @@ import logging
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from ragatouille import RAGPretrainedModel
 from langchain.chains import LLMChain
-from langchain.llms import OpenAI
+from langchain_community.llms import OpenAI
 from langchain import PromptTemplate
+import json
+from concurrent.futures import ThreadPoolExecutor
 
-try:
-    from modules.helpers import get_metadata
-except:
-    from helpers import get_metadata
-
-logger = logging.getLogger(__name__)
+from modules.dataloader.helpers import get_metadata
 
 
 class PDFReader:
@@ -40,8 +37,9 @@ class PDFReader:
 
 
 class FileReader:
-    def __init__(self):
+    def __init__(self, logger):
         self.pdf_reader = PDFReader()
+        self.logger = logger
 
     def extract_text_from_pdf(self, pdf_path):
         text = ""
@@ -61,7 +59,7 @@ class FileReader:
                 temp_file_path = temp_file.name
             return temp_file_path
         else:
-            print("Failed to download PDF from URL:", pdf_url)
+            self.logger.error(f"Failed to download PDF from URL: {pdf_url}")
             return None
 
     def read_pdf(self, temp_file_path: str):
@@ -99,13 +97,18 @@ class FileReader:
         if response.status_code == 200:
             return [Document(page_content=response.text)]
         else:
-            print("Failed to fetch .tex file from URL:", tex_url)
+            self.logger.error(f"Failed to fetch .tex file from URL: {tex_url}")
             return None
 
 
 class ChunkProcessor:
-    def __init__(self, config):
+    def __init__(self, config, logger):
         self.config = config
+        self.logger = logger
+
+        self.document_data = {}
+        self.document_metadata = {}
+        self.document_chunks_full = []
 
         if config["splitter_options"]["use_splitter"]:
             if config["splitter_options"]["split_by_token"]:
@@ -124,7 +127,7 @@ class ChunkProcessor:
                 )
         else:
             self.splitter = None
-        logger.info("ChunkProcessor instance created")
+        self.logger.info("ChunkProcessor instance created")
 
     def remove_delimiters(self, document_chunks: list):
         for chunk in document_chunks:
@@ -139,7 +142,6 @@ class ChunkProcessor:
             del document_chunks[0]
         for _ in range(end):
             document_chunks.pop()
-        logger.info(f"\tNumber of pages after skipping: {len(document_chunks)}")
         return document_chunks
 
     def process_chunks(
@@ -172,122 +174,184 @@ class ChunkProcessor:
 
         return document_chunks
 
-    def get_chunks(self, file_reader, uploaded_files, weblinks):
-        self.document_chunks_full = []
-        self.parent_document_names = []
-        self.child_document_names = []
-        self.documents = []
-        self.document_metadata = []
-
+    def chunk_docs(self, file_reader, uploaded_files, weblinks):
         addl_metadata = get_metadata(
             "https://dl4ds.github.io/sp2024/lectures/",
             "https://dl4ds.github.io/sp2024/schedule/",
         )  # For any additional metadata
 
-        for file_index, file_path in enumerate(uploaded_files):
-            file_name = os.path.basename(file_path)
-            if file_name not in self.parent_document_names:
-                file_type = file_name.split(".")[-1].lower()
-
-                # try:
-                if file_type == "pdf":
-                    documents = file_reader.read_pdf(file_path)
-                elif file_type == "txt":
-                    documents = file_reader.read_txt(file_path)
-                elif file_type == "docx":
-                    documents = file_reader.read_docx(file_path)
-                elif file_type == "srt":
-                    documents = file_reader.read_srt(file_path)
-                elif file_type == "tex":
-                    documents = file_reader.read_tex_from_url(file_path)
-                else:
-                    logger.warning(f"Unsupported file type: {file_type}")
-                    continue
-
-                for doc in documents:
-                    page_num = doc.metadata.get("page", 0)
-                    self.documents.append(doc.page_content)
-                    self.document_metadata.append(
-                        {"source": file_path, "page": page_num}
-                    )
-                    metadata = addl_metadata.get(file_path, {})
-                    self.document_metadata[-1].update(metadata)
-
-                    self.child_document_names.append(f"{file_name}_{page_num}")
-
-                    self.parent_document_names.append(file_name)
-                    if self.config["embedding_options"]["db_option"] not in [
-                        "RAGatouille"
-                    ]:
-                        document_chunks = self.process_chunks(
-                            self.documents[-1],
-                            file_type,
-                            source=file_path,
-                            page=page_num,
-                            metadata=metadata,
-                        )
-                        self.document_chunks_full.extend(document_chunks)
-
-            # except Exception as e:
-            #     logger.error(f"Error processing file {file_name}: {str(e)}")
-
-        self.process_weblinks(file_reader, weblinks)
-
-        logger.info(
+        with ThreadPoolExecutor() as executor:
+            executor.map(
+                self.process_file,
+                uploaded_files,
+                range(len(uploaded_files)),
+                [file_reader] * len(uploaded_files),
+                [addl_metadata] * len(uploaded_files),
+            )
+            executor.map(
+                self.process_weblink,
+                weblinks,
+                range(len(weblinks)),
+                [file_reader] * len(weblinks),
+                [addl_metadata] * len(weblinks),
+            )
+
+        document_names = [
+            f"{file_name}_{page_num}"
+            for file_name, pages in self.document_data.items()
+            for page_num in pages.keys()
+        ]
+        documents = [
+            page for doc in self.document_data.values() for page in doc.values()
+        ]
+        document_metadata = [
+            page for doc in self.document_metadata.values() for page in doc.values()
+        ]
+
+        self.save_document_data()
+
+        self.logger.info(
             f"Total document chunks extracted: {len(self.document_chunks_full)}"
         )
-        return (
-            self.document_chunks_full,
-            self.child_document_names,
-            self.documents,
-            self.document_metadata,
-        )
 
-    def process_weblinks(self, file_reader, weblinks):
-        if weblinks[0] != "":
-            logger.info(f"Splitting weblinks: total of {len(weblinks)}")
-
-            for link_index, link in enumerate(weblinks):
-                if link not in self.parent_document_names:
-                    try:
-                        logger.info(f"\tSplitting link {link_index+1} : {link}")
-                        if "youtube" in link:
-                            documents = file_reader.read_youtube_transcript(link)
-                        else:
-                            documents = file_reader.read_html(link)
-
-                        for doc in documents:
-                            page_num = doc.metadata.get("page", 0)
-                            self.documents.append(doc.page_content)
-                            self.document_metadata.append(
-                                {"source": link, "page": page_num}
-                            )
-                            self.child_document_names.append(f"{link}")
-
-                        self.parent_document_names.append(link)
-                        if self.config["embedding_options"]["db_option"] not in [
-                            "RAGatouille"
-                        ]:
-                            document_chunks = self.process_chunks(
-                                self.documents[-1],
-                                "txt",
-                                source=link,
-                                page=0,
-                                metadata={"source_type": "webpage"},
-                            )
-                            self.document_chunks_full.extend(document_chunks)
-                    except Exception as e:
-                        logger.error(
-                            f"Error splitting link {link_index+1} : {link}: {str(e)}"
-                        )
+        return self.document_chunks_full, document_names, documents, document_metadata
+
+    def process_documents(
+        self, documents, file_path, file_type, metadata_source, addl_metadata
+    ):
+        file_data = {}
+        file_metadata = {}
+
+        for doc in documents:
+            if len(doc.page_content) <= 400:
+                continue
+
+            page_num = doc.metadata.get("page", 0)
+            file_data[page_num] = doc.page_content
+            metadata = (
+                addl_metadata.get(file_path, {})
+                if metadata_source == "file"
+                else {"source": file_path, "page": page_num}
+            )
+            file_metadata[page_num] = metadata
+
+            if self.config["vectorstore"]["db_option"] not in ["RAGatouille"]:
+                document_chunks = self.process_chunks(
+                    doc.page_content,
+                    file_type,
+                    source=file_path,
+                    page=page_num,
+                    metadata=metadata,
+                )
+                self.document_chunks_full.extend(document_chunks)
+
+        self.document_data[file_path] = file_data
+        self.document_metadata[file_path] = file_metadata
+
+    def process_file(self, file_path, file_index, file_reader, addl_metadata):
+        file_name = os.path.basename(file_path)
+        if file_name in self.document_data:
+            return
+
+        file_type = file_name.split(".")[-1].lower()
+        self.logger.info(f"Reading file {file_index + 1}: {file_path}")
+
+        read_methods = {
+            "pdf": file_reader.read_pdf,
+            "txt": file_reader.read_txt,
+            "docx": file_reader.read_docx,
+            "srt": file_reader.read_srt,
+            "tex": file_reader.read_tex_from_url,
+        }
+        if file_type not in read_methods:
+            self.logger.warning(f"Unsupported file type: {file_type}")
+            return
+
+        try:
+            documents = read_methods[file_type](file_path)
+            self.process_documents(
+                documents, file_path, file_type, "file", addl_metadata
+            )
+        except Exception as e:
+            self.logger.error(f"Error processing file {file_name}: {str(e)}")
+
+    def process_weblink(self, link, link_index, file_reader, addl_metadata):
+        if link in self.document_data:
+            return
+
+        self.logger.info(f"Reading link {link_index + 1} : {link}")
+
+        try:
+            if "youtube" in link:
+                documents = file_reader.read_youtube_transcript(link)
+            else:
+                documents = file_reader.read_html(link)
+
+            self.process_documents(documents, link, "txt", "link", addl_metadata)
+        except Exception as e:
+            self.logger.error(f"Error Reading link {link_index + 1} : {link}: {str(e)}")
+
+    def save_document_data(self):
+        if not os.path.exists(f"{self.config['log_chunk_dir']}/docs"):
+            os.makedirs(f"{self.config['log_chunk_dir']}/docs")
+            self.logger.info(
+                f"Creating directory {self.config['log_chunk_dir']}/docs for document data"
+            )
+        self.logger.info(
+            f"Saving document content to {self.config['log_chunk_dir']}/docs/doc_content.json"
+        )
+        if not os.path.exists(f"{self.config['log_chunk_dir']}/metadata"):
+            os.makedirs(f"{self.config['log_chunk_dir']}/metadata")
+            self.logger.info(
+                f"Creating directory {self.config['log_chunk_dir']}/metadata for document metadata"
+            )
+        self.logger.info(
+            f"Saving document metadata to {self.config['log_chunk_dir']}/metadata/doc_metadata.json"
+        )
+        with open(
+            f"{self.config['log_chunk_dir']}/docs/doc_content.json", "w"
+        ) as json_file:
+            json.dump(self.document_data, json_file, indent=4)
+        with open(
+            f"{self.config['log_chunk_dir']}/metadata/doc_metadata.json", "w"
+        ) as json_file:
+            json.dump(self.document_metadata, json_file, indent=4)
+
+    def load_document_data(self):
+        with open(
+            f"{self.config['log_chunk_dir']}/docs/doc_content.json", "r"
+        ) as json_file:
+            self.document_data = json.load(json_file)
+        with open(
+            f"{self.config['log_chunk_dir']}/metadata/doc_metadata.json", "r"
+        ) as json_file:
+            self.document_metadata = json.load(json_file)
 
 
 class DataLoader:
-    def __init__(self, config):
-        self.file_reader = FileReader()
-        self.chunk_processor = ChunkProcessor(config)
+    def __init__(self, config, logger=None):
+        self.file_reader = FileReader(logger=logger)
+        self.chunk_processor = ChunkProcessor(config, logger=logger)
 
     def get_chunks(self, uploaded_files, weblinks):
-        return self.chunk_processor.get_chunks(
+        return self.chunk_processor.chunk_docs(
             self.file_reader, uploaded_files, weblinks
         )
+
+
+if __name__ == "__main__":
+    import yaml
+
+    logger = logging.getLogger(__name__)
+    logger.setLevel(logging.INFO)
+
+    with open("../code/config.yml", "r") as f:
+        config = yaml.safe_load(f)
+
+    data_loader = DataLoader(config, logger=logger)
+    document_chunks, document_names, documents, document_metadata = (
+        data_loader.get_chunks(
+            [],
+            ["https://dl4ds.github.io/sp2024/"],
+        )
+    )

code/modules/dataloader/helpers.py

@@ -0,0 +1,108 @@
+import requests
+from bs4 import BeautifulSoup
+from tqdm import tqdm
+
+
+def get_urls_from_file(file_path: str):
+    """
+    Function to get urls from a file
+    """
+    with open(file_path, "r") as f:
+        urls = f.readlines()
+    urls = [url.strip() for url in urls]
+    return urls
+
+
+def get_base_url(url):
+    parsed_url = urlparse(url)
+    base_url = f"{parsed_url.scheme}://{parsed_url.netloc}/"
+    return base_url
+
+
+def get_metadata(lectures_url, schedule_url):
+    """
+    Function to get the lecture metadata from the lectures and schedule URLs.
+    """
+    lecture_metadata = {}
+
+    # Get the main lectures page content
+    r_lectures = requests.get(lectures_url)
+    soup_lectures = BeautifulSoup(r_lectures.text, "html.parser")
+
+    # Get the main schedule page content
+    r_schedule = requests.get(schedule_url)
+    soup_schedule = BeautifulSoup(r_schedule.text, "html.parser")
+
+    # Find all lecture blocks
+    lecture_blocks = soup_lectures.find_all("div", class_="lecture-container")
+
+    # Create a mapping from slides link to date
+    date_mapping = {}
+    schedule_rows = soup_schedule.find_all("li", class_="table-row-lecture")
+    for row in schedule_rows:
+        try:
+            date = (
+                row.find("div", {"data-label": "Date"}).get_text(separator=" ").strip()
+            )
+            description_div = row.find("div", {"data-label": "Description"})
+            slides_link_tag = description_div.find("a", title="Download slides")
+            slides_link = slides_link_tag["href"].strip() if slides_link_tag else None
+            slides_link = (
+                f"https://dl4ds.github.io{slides_link}" if slides_link else None
+            )
+            if slides_link:
+                date_mapping[slides_link] = date
+        except Exception as e:
+            print(f"Error processing schedule row: {e}")
+            continue
+
+    for block in lecture_blocks:
+        try:
+            # Extract the lecture title
+            title = block.find("span", style="font-weight: bold;").text.strip()
+
+            # Extract the TL;DR
+            tldr = block.find("strong", text="tl;dr:").next_sibling.strip()
+
+            # Extract the link to the slides
+            slides_link_tag = block.find("a", title="Download slides")
+            slides_link = slides_link_tag["href"].strip() if slides_link_tag else None
+            slides_link = (
+                f"https://dl4ds.github.io{slides_link}" if slides_link else None
+            )
+
+            # Extract the link to the lecture recording
+            recording_link_tag = block.find("a", title="Download lecture recording")
+            recording_link = (
+                recording_link_tag["href"].strip() if recording_link_tag else None
+            )
+
+            # Extract suggested readings or summary if available
+            suggested_readings_tag = block.find("p", text="Suggested Readings:")
+            if suggested_readings_tag:
+                suggested_readings = suggested_readings_tag.find_next_sibling("ul")
+                if suggested_readings:
+                    suggested_readings = suggested_readings.get_text(
+                        separator="\n"
+                    ).strip()
+                else:
+                    suggested_readings = "No specific readings provided."
+            else:
+                suggested_readings = "No specific readings provided."
+
+            # Get the date from the schedule
+            date = date_mapping.get(slides_link, "No date available")
+
+            # Add to the dictionary
+            lecture_metadata[slides_link] = {
+                "date": date,
+                "tldr": tldr,
+                "title": title,
+                "lecture_recording": recording_link,
+                "suggested_readings": suggested_readings,
+            }
+        except Exception as e:
+            print(f"Error processing block: {e}")
+            continue
+
+    return lecture_metadata

code/modules/dataloader/webpage_crawler.py

@@ -0,0 +1,115 @@
+import aiohttp
+from aiohttp import ClientSession
+import asyncio
+import requests
+from bs4 import BeautifulSoup
+from urllib.parse import urlparse, urljoin, urldefrag
+
+class WebpageCrawler:
+    def __init__(self):
+        self.dict_href_links = {}
+
+    async def fetch(self, session: ClientSession, url: str) -> str:
+        async with session.get(url) as response:
+            try:
+                return await response.text()
+            except UnicodeDecodeError:
+                return await response.text(encoding="latin1")
+
+    def url_exists(self, url: str) -> bool:
+        try:
+            response = requests.head(url)
+            return response.status_code == 200
+        except requests.ConnectionError:
+            return False
+
+    async def get_links(self, session: ClientSession, website_link: str, base_url: str):
+        html_data = await self.fetch(session, website_link)
+        soup = BeautifulSoup(html_data, "html.parser")
+        list_links = []
+        for link in soup.find_all("a", href=True):
+            href = link["href"].strip()
+            full_url = urljoin(base_url, href)
+            normalized_url = self.normalize_url(full_url)  # sections removed
+            if (
+                normalized_url not in self.dict_href_links
+                and self.is_child_url(normalized_url, base_url)
+                and self.url_exists(normalized_url)
+            ):
+                self.dict_href_links[normalized_url] = None
+                list_links.append(normalized_url)
+
+        return list_links
+
+    async def get_subpage_links(
+        self, session: ClientSession, urls: list, base_url: str
+    ):
+        tasks = [self.get_links(session, url, base_url) for url in urls]
+        results = await asyncio.gather(*tasks)
+        all_links = [link for sublist in results for link in sublist]
+        return all_links
+
+    async def get_all_pages(self, url: str, base_url: str):
+        async with aiohttp.ClientSession() as session:
+            dict_links = {url: "Not-checked"}
+            counter = None
+            while counter != 0:
+                unchecked_links = [
+                    link
+                    for link, status in dict_links.items()
+                    if status == "Not-checked"
+                ]
+                if not unchecked_links:
+                    break
+                new_links = await self.get_subpage_links(
+                    session, unchecked_links, base_url
+                )
+                for link in unchecked_links:
+                    dict_links[link] = "Checked"
+                    print(f"Checked: {link}")
+                dict_links.update(
+                    {
+                        link: "Not-checked"
+                        for link in new_links
+                        if link not in dict_links
+                    }
+                )
+                counter = len(
+                    [
+                        status
+                        for status in dict_links.values()
+                        if status == "Not-checked"
+                    ]
+                )
+
+            checked_urls = [
+                url for url, status in dict_links.items() if status == "Checked"
+            ]
+            return checked_urls
+
+    def is_webpage(self, url: str) -> bool:
+        try:
+            response = requests.head(url, allow_redirects=True)
+            content_type = response.headers.get("Content-Type", "").lower()
+            return "text/html" in content_type
+        except requests.RequestException:
+            return False
+
+    def clean_url_list(self, urls):
+        files, webpages = [], []
+
+        for url in urls:
+            if self.is_webpage(url):
+                webpages.append(url)
+            else:
+                files.append(url)
+
+        return files, webpages
+
+    def is_child_url(self, url, base_url):
+        return url.startswith(base_url)
+
+    def normalize_url(self, url: str):
+        # Strip the fragment identifier
+        defragged_url, _ = urldefrag(url)
+        return defragged_url

code/modules/retriever/__init__.py

@@ -0,0 +1,2 @@
+from .faiss_retriever import FaissRetriever
+from .chroma_retriever import ChromaRetriever

code/modules/retriever/base.py

@@ -0,0 +1,6 @@
+class BaseRetriever:
+    def __init__(self, config):
+        self.config = config
+
+    def return_retriever(self):
+        raise NotImplementedError

code/modules/retriever/chroma_retriever.py

@@ -0,0 +1,24 @@
+from .helpers import VectorStoreRetrieverScore
+from .base import BaseRetriever
+
+
+class ChromaRetriever(BaseRetriever):
+    def __init__(self):
+        pass
+
+    def return_retriever(self, db, config):
+        retriever = VectorStoreRetrieverScore(
+            vectorstore=db,
+            # search_type="similarity_score_threshold",
+            # search_kwargs={
+            #     "score_threshold": self.config["vectorstore"][
+            #         "score_threshold"
+            #     ],
+            #     "k": self.config["vectorstore"]["search_top_k"],
+            # },
+            search_kwargs={
+                "k": config["vectorstore"]["search_top_k"],
+            },
+        )
+
+        return retriever

code/modules/retriever/faiss_retriever.py

@@ -0,0 +1,23 @@
+from .helpers import VectorStoreRetrieverScore
+from .base import BaseRetriever
+
+
+class FaissRetriever(BaseRetriever):
+    def __init__(self):
+        pass
+
+    def return_retriever(self, db, config):
+        retriever = VectorStoreRetrieverScore(
+            vectorstore=db,
+            # search_type="similarity_score_threshold",
+            # search_kwargs={
+            #     "score_threshold": self.config["vectorstore"][
+            #         "score_threshold"
+            #     ],
+            #     "k": self.config["vectorstore"]["search_top_k"],
+            # },
+            search_kwargs={
+                "k": config["vectorstore"]["search_top_k"],
+            },
+        )
+        return retriever

code/modules/retriever/helpers.py

@@ -0,0 +1,39 @@
+from langchain.schema.vectorstore import VectorStoreRetriever
+from langchain.callbacks.manager import CallbackManagerForRetrieverRun
+from langchain.schema.document import Document
+from langchain_core.callbacks import AsyncCallbackManagerForRetrieverRun
+from typing import List
+
+
+class VectorStoreRetrieverScore(VectorStoreRetriever):
+
+    # See https://github.com/langchain-ai/langchain/blob/61dd92f8215daef3d9cf1734b0d1f8c70c1571c3/libs/langchain/langchain/vectorstores/base.py#L500
+    def _get_relevant_documents(
+        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> List[Document]:
+        docs_and_similarities = (
+            self.vectorstore.similarity_search_with_relevance_scores(
+                query, **self.search_kwargs
+            )
+        )
+        # Make the score part of the document metadata
+        for doc, similarity in docs_and_similarities:
+            doc.metadata["score"] = similarity
+
+        docs = [doc for doc, _ in docs_and_similarities]
+        return docs
+
+    async def _aget_relevant_documents(
+        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
+    ) -> List[Document]:
+        docs_and_similarities = (
+            self.vectorstore.similarity_search_with_relevance_scores(
+                query, **self.search_kwargs
+            )
+        )
+        # Make the score part of the document metadata
+        for doc, similarity in docs_and_similarities:
+            doc.metadata["score"] = similarity
+
+        docs = [doc for doc, _ in docs_and_similarities]
+        return docs

code/modules/vectorstore/__init__.py

@@ -0,0 +1,2 @@
+from .base import VectorStoreBase
+from .faiss import FAISS

code/modules/vectorstore/base.py

@@ -0,0 +1,18 @@
+class VectorStoreBase:
+    def __init__(self, config):
+        self.config = config
+
+    def _init_vector_db(self):
+        raise NotImplementedError
+
+    def create_database(self, database_name):
+        raise NotImplementedError
+
+    def load_database(self, database_name):
+        raise NotImplementedError
+
+    def as_retriever(self):
+        raise NotImplementedError
+
+    def __str__(self):
+        return self.__class__.__name__

code/modules/vectorstore/chroma.py

@@ -0,0 +1,41 @@
+from langchain_community.vectorstores import Chroma
+from modules.vectorstore.base import VectorStoreBase
+import os
+
+
+class ChromaVectorStore(VectorStoreBase):
+    def __init__(self, config):
+        self.config = config
+        self._init_vector_db()
+
+    def _init_vector_db(self):
+        self.chroma = Chroma()
+
+    def create_database(self, document_chunks, embedding_model):
+        self.vectorstore = self.chroma.from_documents(
+            documents=document_chunks,
+            embedding=embedding_model,
+            persist_directory=os.path.join(
+                self.config["vectorstore"]["db_path"],
+                "db_"
+                + self.config["vectorstore"]["db_option"]
+                + "_"
+                + self.config["vectorstore"]["model"],
+            ),
+        )
+
+    def load_database(self, embedding_model):
+        self.vectorstore = Chroma(
+            persist_directory=os.path.join(
+                self.config["vectorstore"]["db_path"],
+                "db_"
+                + self.config["vectorstore"]["db_option"]
+                + "_"
+                + self.config["vectorstore"]["model"],
+            ),
+            embedding_function=embedding_model,
+        )
+        return self.vectorstore
+
+    def as_retriever(self):
+        return self.vectorstore.as_retriever()

code/modules/{embedding_model_loader.py → vectorstore/embedding_model_loader.py}

@@ -2,10 +2,7 @@ from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.embeddings import LlamaCppEmbeddings
 
-try:
-    from modules.constants import *
-except:
-    from constants import *
+from modules.config.constants import *
 import os
 
 
@@ -14,19 +11,19 @@ class EmbeddingModelLoader:
         self.config = config
 
     def load_embedding_model(self):
-        if self.config["embedding_options"]["model"] in ["text-embedding-ada-002"]:
+        if self.config["vectorstore"]["model"] in ["text-embedding-ada-002"]:
             embedding_model = OpenAIEmbeddings(
                 deployment="SL-document_embedder",
-                model=self.config["embedding_options"]["model"],
+                model=self.config["vectorestore"]["model"],
                 show_progress_bar=True,
                 openai_api_key=OPENAI_API_KEY,
                 disallowed_special=(),
             )
         else:
             embedding_model = HuggingFaceEmbeddings(
-                model_name=self.config["embedding_options"]["model"],
+                model_name=self.config["vectorstore"]["model"],
                 model_kwargs={
-                    "device": "cpu",
+                    "device": f"{self.config['device']}",
                     "token": f"{HUGGINGFACE_TOKEN}",
                     "trust_remote_code": True,
                 },

code/modules/vectorstore/faiss.py

@@ -0,0 +1,45 @@
+from langchain_community.vectorstores import FAISS
+from modules.vectorstore.base import VectorStoreBase
+import os
+
+
+class FaissVectorStore(VectorStoreBase):
+    def __init__(self, config):
+        self.config = config
+        self._init_vector_db()
+
+    def _init_vector_db(self):
+        self.faiss = FAISS(
+            embedding_function=None, index=0, index_to_docstore_id={}, docstore={}
+        )
+
+    def create_database(self, document_chunks, embedding_model):
+        self.vectorstore = self.faiss.from_documents(
+            documents=document_chunks, embedding=embedding_model
+        )
+        self.vectorstore.save_local(
+            os.path.join(
+                self.config["vectorstore"]["db_path"],
+                "db_"
+                + self.config["vectorstore"]["db_option"]
+                + "_"
+                + self.config["vectorstore"]["model"],
+            )
+        )
+
+    def load_database(self, embedding_model):
+        self.vectorstore = self.faiss.load_local(
+            os.path.join(
+                self.config["vectorstore"]["db_path"],
+                "db_"
+                + self.config["vectorstore"]["db_option"]
+                + "_"
+                + self.config["vectorstore"]["model"],
+            ),
+            embedding_model,
+            allow_dangerous_deserialization=True,
+        )
+        return self.vectorstore
+
+    def as_retriever(self):
+        return self.vectorstore.as_retriever()

code/modules/{vector_db.py → vectorstore/store_manager.py}

@@ -1,72 +1,27 @@
+from modules.vectorstore.faiss import FaissVectorStore
+from modules.vectorstore.chroma import ChromaVectorStore
+from modules.vectorstore.helpers import *
+from modules.dataloader.webpage_crawler import WebpageCrawler
+from modules.dataloader.data_loader import DataLoader
+from modules.dataloader.helpers import *
+from modules.vectorstore.embedding_model_loader import EmbeddingModelLoader
 import logging
 import os
-import yaml
-from langchain_community.vectorstores import FAISS, Chroma
-from langchain.schema.vectorstore import VectorStoreRetriever
-from langchain.callbacks.manager import CallbackManagerForRetrieverRun
-from langchain.schema.document import Document
-from langchain_core.callbacks import AsyncCallbackManagerForRetrieverRun
-from ragatouille import RAGPretrainedModel
-
-try:
-    from modules.embedding_model_loader import EmbeddingModelLoader
-    from modules.data_loader import DataLoader
-    from modules.constants import *
-    from modules.helpers import *
-except:
-    from embedding_model_loader import EmbeddingModelLoader
-    from data_loader import DataLoader
-    from constants import *
-    from helpers import *
-
-from typing import List
-
-
-class VectorDBScore(VectorStoreRetriever):
-
-    # See https://github.com/langchain-ai/langchain/blob/61dd92f8215daef3d9cf1734b0d1f8c70c1571c3/libs/langchain/langchain/vectorstores/base.py#L500
-    def _get_relevant_documents(
-        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
-    ) -> List[Document]:
-        docs_and_similarities = (
-            self.vectorstore.similarity_search_with_relevance_scores(
-                query, **self.search_kwargs
-            )
-        )
-        # Make the score part of the document metadata
-        for doc, similarity in docs_and_similarities:
-            doc.metadata["score"] = similarity
-
-        docs = [doc for doc, _ in docs_and_similarities]
-        return docs
-
-    async def _aget_relevant_documents(
-        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
-    ) -> List[Document]:
-        docs_and_similarities = (
-            self.vectorstore.similarity_search_with_relevance_scores(
-                query, **self.search_kwargs
-            )
-        )
-        # Make the score part of the document metadata
-        for doc, similarity in docs_and_similarities:
-            doc.metadata["score"] = similarity
-
-        docs = [doc for doc, _ in docs_and_similarities]
-        return docs
+import time
+import asyncio
 
 
-class VectorDB:
+class VectorStoreManager:
     def __init__(self, config, logger=None):
         self.config = config
-        self.db_option = config["embedding_options"]["db_option"]
+        self.db_option = config["vectorstore"]["db_option"]
         self.document_names = None
-        self.webpage_crawler = WebpageCrawler()
 
         # Set up logging to both console and a file
         if logger is None:
             self.logger = logging.getLogger(__name__)
             self.logger.setLevel(logging.INFO)
+            self.logger.propagate = False
 
             # Console Handler
             console_handler = logging.StreamHandler()
@@ -75,8 +30,13 @@ class VectorDB:
             console_handler.setFormatter(formatter)
             self.logger.addHandler(console_handler)
 
+            # Ensure log directory exists
+            log_directory = self.config["log_dir"]
+            if not os.path.exists(log_directory):
+                os.makedirs(log_directory)
+
             # File Handler
-            log_file_path = "vector_db.log"  # Change this to your desired log file path
+            log_file_path = f"{log_directory}/vector_db.log"  # Change this to your desired log file path
             file_handler = logging.FileHandler(log_file_path, mode="w")
             file_handler.setLevel(logging.INFO)
             file_handler.setFormatter(formatter)
@@ -84,16 +44,18 @@ class VectorDB:
         else:
             self.logger = logger
 
+        self.webpage_crawler = WebpageCrawler()
+
         self.logger.info("VectorDB instance instantiated")
 
     def load_files(self):
-        files = os.listdir(self.config["embedding_options"]["data_path"])
+        files = os.listdir(self.config["vectorstore"]["data_path"])
         files = [
-            os.path.join(self.config["embedding_options"]["data_path"], file)
+            os.path.join(self.config["vectorstore"]["data_path"], file)
             for file in files
         ]
-        urls = get_urls_from_file(self.config["embedding_options"]["url_file_path"])
-        if self.config["embedding_options"]["expand_urls"]:
+        urls = get_urls_from_file(self.config["vectorstore"]["url_file_path"])
+        if self.config["vectorstore"]["expand_urls"]:
             all_urls = []
             for url in urls:
                 loop = asyncio.get_event_loop()
@@ -109,8 +71,9 @@ class VectorDB:
 
     def create_embedding_model(self):
         self.logger.info("Creating embedding function")
-        self.embedding_model_loader = EmbeddingModelLoader(self.config)
-        self.embedding_model = self.embedding_model_loader.load_embedding_model()
+        embedding_model_loader = EmbeddingModelLoader(self.config)
+        embedding_model = embedding_model_loader.load_embedding_model()
+        return embedding_model
 
     def initialize_database(
         self,
@@ -120,107 +83,153 @@ class VectorDB:
         document_metadata: list,
     ):
         if self.db_option in ["FAISS", "Chroma"]:
-            self.create_embedding_model()
-        # Track token usage
+            self.embedding_model = self.create_embedding_model()
+
         self.logger.info("Initializing vector_db")
         self.logger.info("\tUsing {} as db_option".format(self.db_option))
         if self.db_option == "FAISS":
-            self.vector_db = FAISS.from_documents(
-                documents=document_chunks, embedding=self.embedding_model
-            )
+            self.vector_db = FaissVectorStore(self.config)
+            self.vector_db.create_database(document_chunks, self.embedding_model)
         elif self.db_option == "Chroma":
-            self.vector_db = Chroma.from_documents(
-                documents=document_chunks,
-                embedding=self.embedding_model,
-                persist_directory=os.path.join(
-                    self.config["embedding_options"]["db_path"],
-                    "db_"
-                    + self.config["embedding_options"]["db_option"]
-                    + "_"
-                    + self.config["embedding_options"]["model"],
-                ),
-            )
+            self.vector_db = ChromaVectorStore(self.config)
+            self.vector_db.create_database(document_chunks, self.embedding_model)
         elif self.db_option == "RAGatouille":
             self.RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")
-            index_path = self.RAG.index(
-                index_name="new_idx",
-                collection=documents,
-                document_ids=document_names,
-                document_metadatas=document_metadata,
-            )
-        self.logger.info("Completed initializing vector_db")
+            # index_path = self.RAG.index(
+            #     index_name="new_idx",
+            #     collection=documents,
+            #     document_ids=document_names,
+            #     document_metadatas=document_metadata,
+            # )
+            batch_size = 32
+            for i in range(0, len(documents), batch_size):
+                if i == 0:
+                    self.RAG.index(
+                        index_name="new_idx",
+                        collection=documents[i : i + batch_size],
+                        document_ids=document_names[i : i + batch_size],
+                        document_metadatas=document_metadata[i : i + batch_size],
+                    )
+                else:
+                    self.RAG = RAGPretrainedModel.from_index(
+                        ".ragatouille/colbert/indexes/new_idx"
+                    )
+                    self.RAG.add_to_index(
+                        new_collection=documents[i : i + batch_size],
+                        new_document_ids=document_names[i : i + batch_size],
+                        new_document_metadatas=document_metadata[i : i + batch_size],
+                    )
 
     def create_database(self):
-        data_loader = DataLoader(self.config)
+        start_time = time.time()  # Start time for creating database
+        data_loader = DataLoader(self.config, self.logger)
         self.logger.info("Loading data")
         files, urls = self.load_files()
         files, webpages = self.webpage_crawler.clean_url_list(urls)
-        if "storage/data/urls.txt" in files:
-            files.remove("storage/data/urls.txt")
+        self.logger.info(f"Number of files: {len(files)}")
+        self.logger.info(f"Number of webpages: {len(webpages)}")
+        if f"{self.config['vectorstore']['url_file_path']}" in files:
+            files.remove(f"{self.config['vectorstores']['url_file_path']}")  # cleanup
         document_chunks, document_names, documents, document_metadata = (
             data_loader.get_chunks(files, webpages)
         )
+        num_documents = len(document_chunks)
+        self.logger.info(f"Number of documents in the DB: {num_documents}")
+        metadata_keys = list(document_metadata[0].keys())
+        self.logger.info(f"Metadata keys: {metadata_keys}")
         self.logger.info("Completed loading data")
         self.initialize_database(
             document_chunks, document_names, documents, document_metadata
         )
+        end_time = time.time()  # End time for creating database
+        self.logger.info("Created database")
+        self.logger.info(
+            f"Time taken to create database: {end_time - start_time} seconds"
+        )
 
-    def save_database(self):
-        if self.db_option == "FAISS":
-            self.vector_db.save_local(
-                os.path.join(
-                    self.config["embedding_options"]["db_path"],
-                    "db_"
-                    + self.config["embedding_options"]["db_option"]
-                    + "_"
-                    + self.config["embedding_options"]["model"],
-                )
-            )
-        elif self.db_option == "Chroma":
-            # db is saved in the persist directory during initialization
-            pass
-        elif self.db_option == "RAGatouille":
-            # index is saved during initialization
-            pass
-        self.logger.info("Saved database")
+    # def save_database(self):
+    #     start_time = time.time()  # Start time for saving database
+    #     if self.db_option == "FAISS":
+    #         self.vector_db.save_local(
+    #             os.path.join(
+    #                 self.config["vectorstore"]["db_path"],
+    #                 "db_"
+    #                 + self.config["vectorstore"]["db_option"]
+    #                 + "_"
+    #                 + self.config["vectorstore"]["model"],
+    #             )
+    #         )
+    #     elif self.db_option == "Chroma":
+    #         # db is saved in the persist directory during initialization
+    #         pass
+    #     elif self.db_option == "RAGatouille":
+    #         # index is saved during initialization
+    #         pass
+    #     self.logger.info("Saved database")
+    #     end_time = time.time()  # End time for saving database
+    #     self.logger.info(
+    #         f"Time taken to save database: {end_time - start_time} seconds"
+    #     )
 
     def load_database(self):
-        self.create_embedding_model()
+        start_time = time.time()  # Start time for loading database
+        if self.db_option in ["FAISS", "Chroma"]:
+            self.embedding_model = self.create_embedding_model()
         if self.db_option == "FAISS":
-            self.vector_db = FAISS.load_local(
-                os.path.join(
-                    self.config["embedding_options"]["db_path"],
-                    "db_"
-                    + self.config["embedding_options"]["db_option"]
-                    + "_"
-                    + self.config["embedding_options"]["model"],
-                ),
-                self.embedding_model,
-                allow_dangerous_deserialization=True,
-            )
+            self.vector_db = FaissVectorStore(self.config)
+            self.loaded_vector_db = self.vector_db.load_database(self.embedding_model)
         elif self.db_option == "Chroma":
-            self.vector_db = Chroma(
-                persist_directory=os.path.join(
-                    self.config["embedding_options"]["db_path"],
-                    "db_"
-                    + self.config["embedding_options"]["db_option"]
-                    + "_"
-                    + self.config["embedding_options"]["model"],
-                ),
-                embedding_function=self.embedding_model,
-            )
-        elif self.db_option == "RAGatouille":
-            self.vector_db = RAGPretrainedModel.from_index(
-                ".ragatouille/colbert/indexes/new_idx"
-            )
+            self.vector_db = ChromaVectorStore(self.config)
+            self.loaded_vector_db = self.vector_db.load_database(self.embedding_model)
+        # if self.db_option == "FAISS":
+        #     self.vector_db = FAISS.load_local(
+        #         os.path.join(
+        #             self.config["vectorstore"]["db_path"],
+        #             "db_"
+        #             + self.config["vectorstore"]["db_option"]
+        #             + "_"
+        #             + self.config["vectorstore"]["model"],
+        #         ),
+        #         self.embedding_model,
+        #         allow_dangerous_deserialization=True,
+        #     )
+        # elif self.db_option == "Chroma":
+        #     self.vector_db = Chroma(
+        #         persist_directory=os.path.join(
+        #             self.config["embedding_options"]["db_path"],
+        #             "db_"
+        #             + self.config["embedding_options"]["db_option"]
+        #             + "_"
+        #             + self.config["embedding_options"]["model"],
+        #         ),
+        #         embedding_function=self.embedding_model,
+        #     )
+        # elif self.db_option == "RAGatouille":
+        #     self.vector_db = RAGPretrainedModel.from_index(
+        #         ".ragatouille/colbert/indexes/new_idx"
+        #     )
+        end_time = time.time()  # End time for loading database
+        self.logger.info(
+            f"Time taken to load database: {end_time - start_time} seconds"
+        )
         self.logger.info("Loaded database")
-        return self.vector_db
+        return self.loaded_vector_db
 
 
 if __name__ == "__main__":
-    with open("code/config.yml", "r") as f:
+    import yaml
+
+    with open("modules/config/config.yml", "r") as f:
         config = yaml.safe_load(f)
     print(config)
-    vector_db = VectorDB(config)
+    print(f"Trying to create database with config: {config}")
+    vector_db = VectorStoreManager(config)
     vector_db.create_database()
-    vector_db.save_database()
+    print("Created database")
+
+    print(f"Trying to load the database")
+    vector_db = VectorStoreManager(config)
+    vector_db.load_database()
+    print("Loaded database")
+
+    print(f"View the logs at {config['log_dir']}/vector_db.log")