Merge pull request #28 from DL4DS/dev_branch

.github/workflows/push_to_hf_space_prototype.yml

@@ -0,0 +1,20 @@
+name: Push Prototype to HuggingFace
+
+on:
+  pull_request:
+    branches:
+    - dev_branch
+
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Deploy Prototype to HuggingFace
+      uses: nateraw/[email protected]
+      with:
+        github_repo_id: DL4DS/dl4ds_tutor
+        huggingface_repo_id: dl4ds/tutor_dev
+        repo_type: space
+        space_sdk: static
+        hf_token: ${{ secrets.HF_TOKEN }}

.gitignore

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

Dockerfile

@@ -1,18 +1,17 @@
-FROM python:3.9
+FROM python:3.11
 
 WORKDIR /code
 
 COPY ./requirements.txt /code/requirements.txt
 
-RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+RUN pip install --upgrade pip
 
-RUN pip install --no-cache-dir transformers==4.36.2 torch==2.1.2
-
-RUN pip install --upgrade --force-reinstall --no-cache-dir llama-cpp-python==0.2.32
+RUN pip install --no-cache-dir -r /code/requirements.txt
 
 COPY . /code
 
-RUN ls -R
+# List the contents of the /code directory to verify files are copied correctly
+RUN ls -R /code
 
 # Change permissions to allow writing to the directory
 RUN chmod -R 777 /code
@@ -23,7 +22,10 @@ RUN mkdir /code/logs && chmod 777 /code/logs
 # Create a cache directory within the application's working directory
 RUN mkdir /.cache && chmod -R 777 /.cache
 
+WORKDIR /code/code
+
 RUN --mount=type=secret,id=HUGGINGFACEHUB_API_TOKEN,mode=0444,required=true 
 RUN --mount=type=secret,id=OPENAI_API_KEY,mode=0444,required=true 
 
-CMD python code/modules/vector_db.py && chainlit run code/main.py --host 0.0.0.0 --port 7860
+# Default command to run the application
+CMD ["sh", "-c", "python -m modules.vectorstore.store_manager && chainlit run main.py --host 0.0.0.0 --port 7860"]

Dockerfile.dev

@@ -0,0 +1,31 @@
+FROM python:3.11
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --upgrade pip
+
+RUN pip install --no-cache-dir -r /code/requirements.txt
+
+COPY . /code
+
+# List the contents of the /code directory to verify files are copied correctly
+RUN ls -R /code
+
+# Change permissions to allow writing to the directory
+RUN chmod -R 777 /code
+
+# Create a logs directory and set permissions
+RUN mkdir /code/logs && chmod 777 /code/logs
+
+# Create a cache directory within the application's working directory
+RUN mkdir /.cache && chmod -R 777 /.cache
+
+WORKDIR /code/code
+
+# Expose the port the app runs on
+EXPOSE 8051
+
+# Default command to run the application
+CMD ["sh", "-c", "python -m modules.vectorstore.store_manager && chainlit run main.py --host 0.0.0.0 --port 8051"]

README.md

@@ -1,35 +1,84 @@
----
-title: Dl4ds Tutor
-emoji: 🏃
-colorFrom: green
-colorTo: red
-sdk: docker
-pinned: false
-hf_oauth: true
----
+# DL4DS Tutor 🏃
 
-DL4DS Tutor - DS598
-===========
+Check out the configuration reference at [Hugging Face Spaces Config Reference](https://huggingface.co/docs/hub/spaces-config-reference).
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+You can find an implementation of the Tutor at [DL4DS Tutor on Hugging Face](https://dl4ds-dl4ds-tutor.hf.space/), which is hosted on Hugging Face [here](https://huggingface.co/spaces/dl4ds/dl4ds_tutor).
 
-You can find an implementation of the Tutor at https://dl4ds-dl4ds-tutor.hf.space/, which is hosted on Hugging Face [here](https://huggingface.co/spaces/dl4ds/dl4ds_tutor)
+## Running Locally
 
-To run locally, 
+1. **Clone the Repository**
+   ```bash
+   git clone https://github.com/DL4DS/dl4ds_tutor
+   ```
 
-Clone the repository from: https://github.com/DL4DS/dl4ds_tutor    
+2. **Put your data under the `storage/data` directory**
+   - Add URLs in the `urls.txt` file.
+   - Add other PDF files in the `storage/data` directory.
 
-Put your data under the `storage/data` directory. Note: You can add urls in the urls.txt file, and other pdf files in the `storage/data` directory.    
+3. **To test Data Loading (Optional)**
+   ```bash
+   cd code
+   python -m modules.dataloader.data_loader
+   ```
 
-To create the Vector Database, run the following command:   
-```python code/modules/vector_db.py```    
-(Note: You would need to run the above when you add new data to the `storage/data` directory, or if the ``storage/data/urls.txt`` file is updated. Or you can set ``["embedding_options"]["embedd_files"]`` to True in the `code/config.yaml` file, which would embed files from the storage directory everytime you run the below chainlit command.)
+4. **Create the Vector Database**
+   ```bash
+   cd code
+   python -m modules.vectorstore.store_manager
+   ```
+   - Note: You need to run the above command when you add new data to the `storage/data` directory, or if the `storage/data/urls.txt` file is updated.
+   - Alternatively, you can set `["vectorstore"]["embedd_files"]` to `True` in the `code/modules/config/config.yaml` file, which will embed files from the storage directory every time you run the below chainlit command.
 
-To run the chainlit app, run the following command:   
-```chainlit run code/main.py```
+5. **Run the Chainlit App**
+   ```bash
+   chainlit run main.py
+   ```
 
 See the [docs](https://github.com/DL4DS/dl4ds_tutor/tree/main/docs) for more information.
 
+## File Structure
+
+```plaintext
+code/
+ ├── modules
+ │   ├── chat                # Contains the chatbot implementation
+ │   ├── chat_processor      # Contains the implementation to process and log the conversations
+ │   ├── config              # Contains the configuration files
+ │   ├── dataloader          # Contains the implementation to load the data from the storage directory
+ │   ├── retriever           # Contains the implementation to create the retriever
+ │   └── vectorstore         # Contains the implementation to create the vector database
+ ├── public
+ │   ├── logo_dark.png       # Dark theme logo
+ │   ├── logo_light.png      # Light theme logo
+ │   └── test.css            # Custom CSS file
+ └── main.py
+
+ 
+docs/                        # Contains the documentation to the codebase and methods used
+
+storage/
+ ├── data                    # Store files and URLs here
+ ├── logs                    # Logs directory, includes logs on vector DB creation, tutor logs, and chunks logged in JSON files
+ └── models                  # Local LLMs are loaded from here
+
+vectorstores/                # Stores the created vector databases
+
+.env                         # This needs to be created, store the API keys here
+```
+- `code/modules/vectorstore/vectorstore.py`: Instantiates the `VectorStore` class to create the vector database.
+- `code/modules/vectorstore/store_manager.py`: Instantiates the `VectorStoreManager:` class to manage the vector database, and all associated methods.
+- `code/modules/retriever/retriever.py`: Instantiates the `Retriever` class to create the retriever.
+
+
+## Docker 
+
+The HuggingFace Space is built using the `Dockerfile` in the repository. To run it locally, use the `Dockerfile.dev` file.
+
+```bash
+docker build --tag dev  -f Dockerfile.dev .
+docker run -it --rm -p 8051:8051 dev
+```
+
 ## Contributing
 
-Please create an issue if you have any suggestions or improvements, and start working on it by creating a branch and by making a pull request to the main branch.
+Please create an issue if you have any suggestions or improvements, and start working on it by creating a branch and by making a pull request to the main branch.

{.chainlit → code/.chainlit}/config.toml

@@ -2,6 +2,7 @@
 # Whether to enable telemetry (default: true). No personal data is collected.
 enable_telemetry = true
 
+
 # List of environment variables to be provided by each user to use the app.
 user_env = []
 
@@ -11,74 +12,104 @@ session_timeout = 3600
 # Enable third parties caching (e.g LangChain cache)
 cache = false
 
+# Authorized origins
+allow_origins = ["*"]
+
 # Follow symlink for asset mount (see https://github.com/Chainlit/chainlit/issues/317)
 # follow_symlink = false
 
 [features]
-# Show the prompt playground
-prompt_playground = true
-
 # Process and display HTML in messages. This can be a security risk (see https://stackoverflow.com/questions/19603097/why-is-it-dangerous-to-render-user-generated-html-or-javascript)
 unsafe_allow_html = false
 
 # Process and display mathematical expressions. This can clash with "$" characters in messages.
 latex = false
 
-# Authorize users to upload files with messages
-multi_modal = true
-
-# Allows user to use speech to text
-[features.speech_to_text]
-    enabled = false
-    # See all languages here https://github.com/JamesBrill/react-speech-recognition/blob/HEAD/docs/API.md#language-string
-    # language = "en-US"
+# Automatically tag threads with the current chat profile (if a chat profile is used)
+auto_tag_thread = true
+
+# Authorize users to spontaneously upload files with messages
+[features.spontaneous_file_upload]
+    enabled = true
+    accept = ["*/*"]
+    max_files = 20
+    max_size_mb = 500
+
+[features.audio]
+    # Threshold for audio recording
+    min_decibels = -45
+    # Delay for the user to start speaking in MS
+    initial_silence_timeout = 3000
+    # Delay for the user to continue speaking in MS. If the user stops speaking for this duration, the recording will stop.
+    silence_timeout = 1500
+    # Above this duration (MS), the recording will forcefully stop.
+    max_duration = 15000
+    # Duration of the audio chunks in MS
+    chunk_duration = 1000
+    # Sample rate of the audio
+    sample_rate = 44100
 
 [UI]
-# Name of the app and chatbot.
+# Name of the assistant.
 name = "AI Tutor"
 
-# Show the readme while the conversation is empty.
-show_readme_as_default = true
-
-# Description of the app and chatbot. This is used for HTML tags.
+# Description of the assistant. This is used for HTML tags.
 # description = ""
 
 # Large size content are by default collapsed for a cleaner ui
 default_collapse_content = true
 
-# The default value for the expand messages settings.
-default_expand_messages = false
-
 # Hide the chain of thought details from the user in the UI.
-hide_cot = false
+hide_cot = true
 
 # Link to your github repo. This will add a github button in the UI's header.
-# github = ""
+# github = "https://github.com/DL4DS/dl4ds_tutor"
 
 # Specify a CSS file that can be used to customize the user interface.
 # The CSS file can be served from the public directory or via an external link.
-# custom_css = "/public/test.css"
+custom_css = "/public/test.css"
 
+# Specify a Javascript file that can be used to customize the user interface.
+# The Javascript file can be served from the public directory.
+# custom_js = "/public/test.js"
+
+# Specify a custom font url.
+# custom_font = "https://fonts.googleapis.com/css2?family=Inter:wght@400;500;700&display=swap"
+
+# Specify a custom meta image url.
+custom_meta_image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/f/f5/Boston_University_seal.svg/1200px-Boston_University_seal.svg.png"
+
+# Specify a custom build directory for the frontend.
+# This can be used to customize the frontend code.
+# Be careful: If this is a relative path, it should not start with a slash.
+# custom_build = "./public/build"
+
+[UI.theme]
+    default = "light"
+    #layout = "wide"
+    #font_family = "Inter, sans-serif"
 # Override default MUI light theme. (Check theme.ts)
 [UI.theme.light]
-    background = "#FFFFFF"
+    background = "#FAFAFA"
     paper = "#FFFFFF"
 
     [UI.theme.light.primary]
-        main = "#000000"
-        dark = "#000000"
-        light = "#FFE7EB"
-
+        main = "#b22222"  # Brighter shade of red
+        dark = "#8b0000"  # Darker shade of the brighter red
+        light = "#ff6347"  # Lighter shade of the brighter red
+    [UI.theme.light.text]
+        primary = "#212121"
+        secondary = "#616161"
 # Override default MUI dark theme. (Check theme.ts)
 [UI.theme.dark]
-    #background = "#FAFAFA"
-    #paper = "#FFFFFF"
+    background = "#1C1C1C" # Slightly lighter dark background color
+    paper = "#2A2A2A"      # Slightly lighter dark paper color
 
     [UI.theme.dark.primary]
-        #main = "#F80061"
-        #dark = "#980039"
-        #light = "#FFE7EB"
+        main = "#89CFF0"    # Primary color
+        dark = "#3700B3"    # Dark variant of primary color
+        light = "#CFBCFF"   # Lighter variant of primary color
 
 
 [meta]
-generated_by = "0.7.700"
+generated_by = "1.1.302"

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 import PromptTemplate
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import FAISS
+from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
+from langchain_core.prompts import PromptTemplate
+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,37 +10,54 @@ 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
-
+from modules.chat.llm_tutor import LLMTutor
+from modules.config.constants import *
+from modules.chat.helpers import get_sources
+from modules.chat_processor.chat_processor import ChatProcessor
 
+global logger
+# Initialize logger
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
+formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
 
 # Console Handler
 console_handler = logging.StreamHandler()
 console_handler.setLevel(logging.INFO)
-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)
+
+@cl.set_starters
+async def set_starters():
+    return [
+        cl.Starter(
+            label="recording on CNNs?",
+            message="Where can I find the recording for the lecture on Transfromers?",
+            icon="/public/adv-screen-recorder-svgrepo-com.svg",
+        ),
+        cl.Starter(
+            label="where's the slides?",
+            message="When are the lectures? I can't find the schedule.",
+            icon="/public/alarmy-svgrepo-com.svg",
+        ),
+        cl.Starter(
+            label="Due Date?",
+            message="When is the final project due?",
+            icon="/public/calendar-samsung-17-svgrepo-com.svg",
+        ),
+        cl.Starter(
+            label="Explain backprop.",
+            message="I didnt understand the math behind backprop, could you explain it?",
+            icon="/public/acastusphoton-svgrepo-com.svg",
+        ),
+    ]
 
 
 # Adding option to select the chat profile
 @cl.set_chat_profiles
 async def chat_profile():
     return [
-        cl.ChatProfile(
-            name="Llama",
-            markdown_description="Use the local LLM: **Tiny Llama**.",
-        ),
         # cl.ChatProfile(
         #     name="Mistral",
         #     markdown_description="Use the local LLM: **Mistral**.",
@@ -54,6 +70,10 @@ async def chat_profile():
             name="gpt-4",
             markdown_description="Use OpenAI API for **gpt-4**.",
         ),
+        cl.ChatProfile(
+            name="Llama",
+            markdown_description="Use the local LLM: **Tiny Llama**.",
+        ),
     ]
 
 
@@ -66,12 +86,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,32 +127,50 @@ 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}...")
-    await msg.send()
-    msg.content = f"Hey, What Can I Help You With?\n\nYou can me ask me questions about the course logistics, course content, about the final project, or anything else!"
-    await msg.update()
+    # msg = cl.Message(content=f"Starting the bot {chat_profile}...")
+    # await msg.send()
+    # msg.content = opening_message
+    # await msg.update()
 
+    tags = [chat_profile, config["vectorstore"]["db_option"]]
+    chat_processor = ChatProcessor(config, tags=tags)
     cl.user_session.set("chain", chain)
+    cl.user_session.set("counter", 0)
+    cl.user_session.set("chat_processor", chat_processor)
+
+
+@cl.on_chat_end
+async def on_chat_end():
+    await cl.Message(content="Sorry, I have to go now. Goodbye!").send()
 
 
 @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}")
+
+    counter = cl.user_session.get("counter")
+    counter += 1
+    cl.user_session.set("counter", counter)
+
+    # if counter >= 3:  # Ensure the counter condition is checked
+    #     await cl.Message(content="Your credits are up!").send()
+    #     await on_chat_end()  # Call the on_chat_end function to handle the end of the chat
+    #     return  # Exit the function to stop further processing
+    # else:
+
+    cb = cl.AsyncLangchainCallbackHandler()  # TODO: fix streaming here
+    cb.answer_reached = True
+
+    processor = cl.user_session.get("chat_processor")
+    res = await processor.rag(message.content, chain, cb)
     try:
         answer = res["answer"]
     except:
         answer = res["result"]
-    print(f"answer: {answer}")
 
-    answer_with_sources, source_elements = get_sources(res, answer)
+    answer_with_sources, source_elements, sources_dict = get_sources(res, answer)
+    processor._process(message.content, answer, sources_dict)
 
     await cl.Message(content=answer_with_sources, elements=source_elements).send()

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/chat/helpers.py

@@ -0,0 +1,104 @@
+from modules.config.constants import *
+import chainlit as cl
+from langchain_core.prompts import PromptTemplate
+
+
+def get_sources(res, answer):
+    source_elements = []
+    source_dict = {}  # Dictionary to store URL elements
+
+    for idx, source in enumerate(res["source_documents"]):
+        source_metadata = source.metadata
+        url = source_metadata.get("source", "N/A")
+        score = source_metadata.get("score", "N/A")
+        page = source_metadata.get("page", 1)
+
+        lecture_tldr = source_metadata.get("tldr", "N/A")
+        lecture_recording = source_metadata.get("lecture_recording", "N/A")
+        suggested_readings = source_metadata.get("suggested_readings", "N/A")
+        date = source_metadata.get("date", "N/A")
+
+        source_type = source_metadata.get("source_type", "N/A")
+
+        url_name = f"{url}_{page}"
+        if url_name not in source_dict:
+            source_dict[url_name] = {
+                "text": source.page_content,
+                "url": url,
+                "score": score,
+                "page": page,
+                "lecture_tldr": lecture_tldr,
+                "lecture_recording": lecture_recording,
+                "suggested_readings": suggested_readings,
+                "date": date,
+                "source_type": source_type,
+            }
+        else:
+            source_dict[url_name]["text"] += f"\n\n{source.page_content}"
+
+    # First, display the answer
+    full_answer = "**Answer:**\n"
+    full_answer += answer
+
+    # Then, display the sources
+    full_answer += "\n\n**Sources:**\n"
+    for idx, (url_name, source_data) in enumerate(source_dict.items()):
+        full_answer += f"\nSource {idx + 1} (Score: {source_data['score']}): {source_data['url']}\n"
+
+        name = f"Source {idx + 1} Text\n"
+        full_answer += name
+        source_elements.append(
+            cl.Text(name=name, content=source_data["text"], display="side")
+        )
+
+        # Add a PDF element if the source is a PDF file
+        if source_data["url"].lower().endswith(".pdf"):
+            name = f"Source {idx + 1} PDF\n"
+            full_answer += name
+            pdf_url = f"{source_data['url']}#page={source_data['page']+1}"
+            source_elements.append(cl.Pdf(name=name, url=pdf_url, display="side"))
+
+    full_answer += "\n**Metadata:**\n"
+    for idx, (url_name, source_data) in enumerate(source_dict.items()):
+        full_answer += f"\nSource {idx + 1} Metadata:\n"
+        source_elements.append(
+            cl.Text(
+                name=f"Source {idx + 1} Metadata",
+                content=f"Source: {source_data['url']}\n"
+                f"Page: {source_data['page']}\n"
+                f"Type: {source_data['source_type']}\n"
+                f"Date: {source_data['date']}\n"
+                f"TL;DR: {source_data['lecture_tldr']}\n"
+                f"Lecture Recording: {source_data['lecture_recording']}\n"
+                f"Suggested Readings: {source_data['suggested_readings']}\n",
+                display="side",
+            )
+        )
+
+    return full_answer, source_elements, source_dict
+
+
+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/chat/llm_tutor.py

@@ -0,0 +1,211 @@
+from langchain.chains import RetrievalQA, ConversationalRetrievalChain
+from langchain.memory import (
+    ConversationBufferWindowMemory,
+    ConversationSummaryBufferMemory,
+)
+from langchain.chains.conversational_retrieval.prompts import QA_PROMPT
+import os
+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.retriever import Retriever
+
+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],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
+        _run_manager = run_manager or AsyncCallbackManagerForChainRun.get_noop_manager()
+        question = inputs["question"]
+        get_chat_history = self._get_chat_history
+        chat_history_str = get_chat_history(inputs["chat_history"])
+        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
+            # )
+            system = (
+                "You are someone that rephrases statements. Rephrase the student's question to add context from their chat history if relevant, ensuring it remains from the student's point of view. "
+                "Incorporate relevant details from the chat history to make the question clearer and more specific."
+                "Do not change the meaning of the original statement, and maintain the student's tone and perspective. "
+                "If the question is conversational and doesn't require context, do not rephrase it. "
+                "Example: If the student previously asked about backpropagation in the context of deep learning and now asks 'what is it', rephrase to 'What is backprogatation.'. "
+                "Example: Do not rephrase if the user is asking something specific like 'cool, suggest a project with transformers to use as my final project'"
+                "Chat history: \n{chat_history_str}\n"
+                "Rephrase the following question only if necessary: '{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
+        accepts_run_manager = (
+            "run_manager" in inspect.signature(self._aget_docs).parameters
+        )
+        if accepts_run_manager:
+            docs = await self._aget_docs(new_question, inputs, run_manager=_run_manager)
+        else:
+            docs = await self._aget_docs(new_question, inputs)  # type: ignore[call-arg]
+
+        output: Dict[str, Any] = {}
+        output["original_question"] = question
+        if self.response_if_no_docs_found is not None and len(docs) == 0:
+            output[self.output_key] = self.response_if_no_docs_found
+        else:
+            new_inputs = inputs.copy()
+            if self.rephrase_question:
+                new_inputs["question"] = new_question
+            new_inputs["chat_history"] = chat_history_str
+
+            # Prepare the final prompt with metadata
+            context = "\n\n".join(
+                [
+                    f"Context {idx+1}: \n(Document content: {doc.page_content}\nMetadata: (source_file: {doc.metadata['source'] if 'source' in doc.metadata else 'unknown'}))"
+                    for idx, doc in enumerate(docs)
+                ]
+            )
+            final_prompt = (
+                "You are an AI Tutor for the course DS598, taught by Prof. Thomas Gardos. Answer the user's question using the provided context. Only use the context if it is relevant. The context is ordered by relevance."
+                "If you don't know the answer, do your best without making things up. Keep the conversation flowing naturally. "
+                "Use chat history and context as guides but avoid repeating past responses. Provide links from the source_file metadata. Use the source context that is most relevent."
+                "Speak in a friendly and engaging manner, like talking to a friend. Avoid sounding repetitive or robotic.\n\n"
+                f"Chat History:\n{chat_history_str}\n\n"
+                f"Context:\n{context}\n\n"
+                "Answer the student's question below in a friendly, concise, and engaging manner. Use the context and history only if relevant, otherwise, engage in a free-flowing conversation.\n"
+                f"Student: {question}\n"
+                "AI Tutor:"
+            )
+
+            # new_inputs["input"] = final_prompt
+            new_inputs["question"] = final_prompt
+            # output["final_prompt"] = final_prompt
+
+            answer = await self.combine_docs_chain.arun(
+                input_documents=docs, callbacks=_run_manager.get_child(), **new_inputs
+            )
+            output[self.output_key] = answer
+
+        if self.return_source_documents:
+            output["source_documents"] = docs
+        output["rephrased_question"] = new_question
+        return output
+
+
+class LLMTutor:
+    def __init__(self, config, logger=None):
+        self.config = config
+        self.llm = self.load_llm()
+        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()
+
+    def set_custom_prompt(self):
+        """
+        Prompt template for QA retrieval for each vectorstore
+        """
+        prompt = get_prompt(self.config)
+        # prompt = QA_PROMPT
+
+        return prompt
+
+    # Retrieval QA Chain
+    def retrieval_qa_chain(self, llm, prompt, db):
+
+        retriever = Retriever(self.config)._return_retriever(db)
+
+        if self.config["llm_params"]["use_history"]:
+            memory = ConversationBufferWindowMemory(
+                k=self.config["llm_params"]["memory_window"],
+                memory_key="chat_history",
+                return_messages=True,
+                output_key="answer",
+                max_token_limit=128,
+            )
+            qa_chain = CustomConversationalRetrievalChain.from_llm(
+                llm=llm,
+                chain_type="stuff",
+                retriever=retriever,
+                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(
+                llm=llm,
+                chain_type="stuff",
+                retriever=retriever,
+                return_source_documents=True,
+                chain_type_kwargs={"prompt": prompt},
+            )
+        return qa_chain
+
+    # Loading the model
+    def load_llm(self):
+        chat_model_loader = ChatModelLoader(self.config)
+        llm = chat_model_loader.load_chat_model()
+        return llm
+
+    # QA Model Function
+    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
+        )  # TODO:  PROMPT is overwritten in CustomConversationalRetrievalChain
+
+        return qa
+
+    # output function
+    def final_result(query):
+        qa_result = qa_bot()
+        response = qa_result({"query": query})
+        return response

code/modules/chat_processor/base.py

@@ -0,0 +1,12 @@
+# Template for chat processor classes
+
+
+class ChatProcessorBase:
+    def __init__(self, config):
+        self.config = config
+
+    def process(self, message):
+        """
+        Processes and Logs the message
+        """
+        raise NotImplementedError("process method not implemented")

code/modules/chat_processor/chat_processor.py

@@ -0,0 +1,30 @@
+from modules.chat_processor.literal_ai import LiteralaiChatProcessor
+
+
+class ChatProcessor:
+    def __init__(self, config, tags=None):
+        self.chat_processor_type = config["chat_logging"]["platform"]
+        self.logging = config["chat_logging"]["log_chat"]
+        self.tags = tags
+        if self.logging:
+            self._init_processor()
+
+    def _init_processor(self):
+        if self.chat_processor_type == "literalai":
+            self.processor = LiteralaiChatProcessor(self.tags)
+        else:
+            raise ValueError(
+                f"Chat processor type {self.chat_processor_type} not supported"
+            )
+
+    def _process(self, user_message, assistant_message, source_dict):
+        if self.logging:
+            return self.processor.process(user_message, assistant_message, source_dict)
+        else:
+            pass
+
+    async def rag(self, user_query: str, chain, cb):
+        if self.logging:
+            return await self.processor.rag(user_query, chain, cb)
+        else:
+            return await chain.acall(user_query, callbacks=[cb])

code/modules/chat_processor/literal_ai.py

@@ -0,0 +1,37 @@
+from literalai import LiteralClient
+import os
+from .base import ChatProcessorBase
+
+
+class LiteralaiChatProcessor(ChatProcessorBase):
+    def __init__(self, tags=None):
+        self.literal_client = LiteralClient(api_key=os.getenv("LITERAL_API_KEY"))
+        self.literal_client.reset_context()
+        with self.literal_client.thread(name="TEST") as thread:
+            self.thread_id = thread.id
+            self.thread = thread
+            if tags is not None and type(tags) == list:
+                self.thread.tags = tags
+        print(f"Thread ID: {self.thread}")
+
+    def process(self, user_message, assistant_message, source_dict):
+        with self.literal_client.thread(thread_id=self.thread_id) as thread:
+            self.literal_client.message(
+                content=user_message,
+                type="user_message",
+                name="User",
+            )
+            self.literal_client.message(
+                content=assistant_message,
+                type="assistant_message",
+                name="AI_Tutor",
+            )
+
+    async def rag(self, user_query: str, chain, cb):
+        with self.literal_client.step(
+            type="retrieval", name="RAG", thread_id=self.thread_id
+        ) as step:
+            step.input = {"question": user_query}
+            res = await chain.acall(user_query, callbacks=[cb])
+            step.output = res
+        return res

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

@@ -1,23 +1,42 @@
-embedding_options:
+log_dir: '../storage/logs' # str
+log_chunk_dir: '../storage/logs/chunks' # str
+device: 'cpu' # str [cuda, cpu]
+
+vectorstore:
   embedd_files: False # bool
-  persist_directory: null # str or None
-  data_path: 'storage/data' # str
-  url_file_path: 'storage/data/urls.txt' # str
+  data_path: '../storage/data' # str
+  url_file_path: '../storage/data/urls.txt' # str
   expand_urls: True # bool
-  db_option : 'FAISS' # str
-  db_path : 'vectorstores' # str
+  db_option : 'FAISS' # str [FAISS, Chroma, RAGatouille, RAPTOR]
+  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
+
+  colbert_params:
+    index_name: "new_idx" # str
+
 llm_params: 
-  use_history: False # bool
+  use_history: True # bool
   memory_window: 3 # int
-  llm_loader: 'local_llm' # str [local_llm, openai]
+  llm_loader: 'openai' # str [local_llm, openai]
   openai_params:
-    model: 'gpt-4' # str [gpt-3.5-turbo-1106, gpt-4]
+    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
+    model: 'tiny-llama'
+    temperature: 0.7
+
+chat_logging:
+  log_chat: False # bool
+  platform: 'literalai'
+
 splitter_options:
   use_splitter: True # bool
   split_by_token : True # bool

code/modules/{constants.py → config/constants.py}

@@ -6,7 +6,10 @@ load_dotenv()
 # API Keys - Loaded from the .env file
 
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+LITERAL_API_KEY = os.getenv("LITERAL_API_KEY")
 
+opening_message = f"Hey, What Can I Help You With?\n\nYou can me ask me questions about the course logistics, course content, about the final project, or anything else!"
 
 # Prompt Templates
 
@@ -75,5 +78,5 @@ Question: {question}
 
 # Model Paths
 
-LLAMA_PATH = "storage/models/tinyllama-1.1b-chat-v1.0.Q5_K_M.gguf"
+LLAMA_PATH = "../storage/models/tinyllama-1.1b-chat-v1.0.Q5_K_M.gguf"
 MISTRAL_PATH = "storage/models/mistral-7b-v0.1.Q4_K_M.gguf"

code/modules/data_loader.py

@@ -1,287 +0,0 @@
-import re
-import pysrt
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.document_loaders import (
-    PyMuPDFLoader,
-    Docx2txtLoader,
-    YoutubeLoader,
-    WebBaseLoader,
-    TextLoader,
-)
-from langchain.schema import Document
-import tempfile
-from tempfile import NamedTemporaryFile
-import logging
-import requests
-
-logger = logging.getLogger(__name__)
-
-
-class DataLoader:
-    def __init__(self, config):
-        """
-        Class for handling all data extraction and chunking
-        Inputs:
-            config - dictionary from yaml file, containing all important parameters
-        """
-        self.config = config
-        self.remove_leftover_delimiters = config["splitter_options"][
-            "remove_leftover_delimiters"
-        ]
-
-        # Main list of all documents
-        self.document_chunks_full = []
-        self.document_names = []
-
-        if config["splitter_options"]["use_splitter"]:
-            if config["splitter_options"]["split_by_token"]:
-                self.splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
-                    chunk_size=config["splitter_options"]["chunk_size"],
-                    chunk_overlap=config["splitter_options"]["chunk_overlap"],
-                    separators=config["splitter_options"]["chunk_separators"],
-                    disallowed_special=()
-                )
-            else:
-                self.splitter = RecursiveCharacterTextSplitter(
-                    chunk_size=config["splitter_options"]["chunk_size"],
-                    chunk_overlap=config["splitter_options"]["chunk_overlap"],
-                    separators=config["splitter_options"]["chunk_separators"],
-                    disallowed_special=()
-                )
-        else:
-            self.splitter = None
-        logger.info("InfoLoader instance created")
-
-    def extract_text_from_pdf(self, pdf_path):
-        text = ""
-        with open(pdf_path, "rb") as file:
-            reader = PyPDF2.PdfReader(file)
-            num_pages = len(reader.pages)
-            for page_num in range(num_pages):
-                page = reader.pages[page_num]
-                text += page.extract_text()
-        return text
-
-    def download_pdf_from_url(self, pdf_url):
-        response = requests.get(pdf_url)
-        if response.status_code == 200:
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
-                temp_file.write(response.content)
-                temp_file_path = temp_file.name
-            return temp_file_path
-        else:
-            print("Failed to download PDF from URL:", pdf_url)
-            return None
-
-    def get_chunks(self, uploaded_files, weblinks):
-        # Main list of all documents
-        self.document_chunks_full = []
-        self.document_names = []
-
-        def remove_delimiters(document_chunks: list):
-            """
-            Helper function to remove remaining delimiters in document chunks
-            """
-            for chunk in document_chunks:
-                for delimiter in self.config["splitter_options"][
-                    "delimiters_to_remove"
-                ]:
-                    chunk.page_content = re.sub(delimiter, " ", chunk.page_content)
-            return document_chunks
-
-        def remove_chunks(document_chunks: list):
-            """
-            Helper function to remove any unwanted document chunks after splitting
-            """
-            front = self.config["splitter_options"]["front_chunk_to_remove"]
-            end = self.config["splitter_options"]["last_chunks_to_remove"]
-            # Remove pages
-            for _ in range(front):
-                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 get_pdf_from_url(pdf_url: str):
-            temp_pdf_path = self.download_pdf_from_url(pdf_url)
-            if temp_pdf_path:
-                title, document_chunks = get_pdf(temp_pdf_path, pdf_url)
-                os.remove(temp_pdf_path)
-                return title, document_chunks
-
-        def get_pdf(temp_file_path: str, title: str):
-            """
-            Function to process PDF files
-            """
-            loader = PyMuPDFLoader(
-                temp_file_path
-            )  # This loader preserves more metadata
-
-            if self.splitter:
-                document_chunks = self.splitter.split_documents(loader.load())
-            else:
-                document_chunks = loader.load()
-
-            if "title" in document_chunks[0].metadata.keys():
-                title = document_chunks[0].metadata["title"]
-
-            logger.info(
-                f"\t\tOriginal no. of pages: {document_chunks[0].metadata['total_pages']}"
-            )
-
-            return title, document_chunks
-
-        def get_txt(temp_file_path: str, title: str):
-            """
-            Function to process TXT files
-            """
-            loader = TextLoader(temp_file_path, autodetect_encoding=True)
-
-            if self.splitter:
-                document_chunks = self.splitter.split_documents(loader.load())
-            else:
-                document_chunks = loader.load()
-
-            # Update the metadata
-            for chunk in document_chunks:
-                chunk.metadata["source"] = title
-                chunk.metadata["page"] = "N/A"
-
-            return title, document_chunks
-
-        def get_srt(temp_file_path: str, title: str):
-            """
-            Function to process SRT files
-            """
-            subs = pysrt.open(temp_file_path)
-
-            text = ""
-            for sub in subs:
-                text += sub.text
-            document_chunks = [Document(page_content=text)]
-
-            if self.splitter:
-                document_chunks = self.splitter.split_documents(document_chunks)
-
-            # Update the metadata
-            for chunk in document_chunks:
-                chunk.metadata["source"] = title
-                chunk.metadata["page"] = "N/A"
-
-            return title, document_chunks
-
-        def get_docx(temp_file_path: str, title: str):
-            """
-            Function to process DOCX files
-            """
-            loader = Docx2txtLoader(temp_file_path)
-
-            if self.splitter:
-                document_chunks = self.splitter.split_documents(loader.load())
-            else:
-                document_chunks = loader.load()
-
-            # Update the metadata
-            for chunk in document_chunks:
-                chunk.metadata["source"] = title
-                chunk.metadata["page"] = "N/A"
-
-            return title, document_chunks
-
-        def get_youtube_transcript(url: str):
-            """
-            Function to retrieve youtube transcript and process text
-            """
-            loader = YoutubeLoader.from_youtube_url(
-                url, add_video_info=True, language=["en"], translation="en"
-            )
-
-            if self.splitter:
-                document_chunks = self.splitter.split_documents(loader.load())
-            else:
-                document_chunks = loader.load_and_split()
-
-            # Replace the source with title (for display in st UI later)
-            for chunk in document_chunks:
-                chunk.metadata["source"] = chunk.metadata["title"]
-            logger.info(chunk.metadata["title"])
-
-            return title, document_chunks
-
-        def get_html(url: str):
-            """
-            Function to process websites via HTML files
-            """
-            loader = WebBaseLoader(url)
-
-            if self.splitter:
-                document_chunks = self.splitter.split_documents(loader.load())
-            else:
-                document_chunks = loader.load_and_split()
-
-            title = document_chunks[0].metadata["title"]
-            logger.info(document_chunks[0].metadata)
-
-            return title, document_chunks
-
-        # Handle file by file
-        for file_index, file_path in enumerate(uploaded_files):
-
-            file_name = file_path.split("/")[-1]
-            file_type = file_name.split(".")[-1]
-
-            # Handle different file types
-            if file_type == "pdf":
-                try:
-                    title, document_chunks = get_pdf(file_path, file_name)
-                except:
-                    title, document_chunks = get_pdf_from_url(file_path)
-            elif file_type == "txt":
-                title, document_chunks = get_txt(file_path, file_name)
-            elif file_type == "docx":
-                title, document_chunks = get_docx(file_path, file_name)
-            elif file_type == "srt":
-                title, document_chunks = get_srt(file_path, file_name)
-
-            # Additional wrangling - Remove leftover delimiters and any specified chunks
-            if self.remove_leftover_delimiters:
-                document_chunks = remove_delimiters(document_chunks)
-            if self.config["splitter_options"]["remove_chunks"]:
-                document_chunks = remove_chunks(document_chunks)
-
-            logger.info(f"\t\tExtracted no. of chunks: {len(document_chunks)} from {file_name}")
-            self.document_names.append(title)
-            self.document_chunks_full.extend(document_chunks)
-
-        # Handle youtube links:
-        if weblinks[0] != "":
-            logger.info(f"Splitting weblinks: total of {len(weblinks)}")
-
-            # Handle link by link
-            for link_index, link in enumerate(weblinks):
-                try:
-                    logger.info(f"\tSplitting link {link_index+1} : {link}")
-                    if "youtube" in link:
-                        title, document_chunks = get_youtube_transcript(link)
-                    else:
-                        title, document_chunks = get_html(link)
-
-                    # Additional wrangling - Remove leftover delimiters and any specified chunks
-                    if self.remove_leftover_delimiters:
-                        document_chunks = remove_delimiters(document_chunks)
-                    if self.config["splitter_options"]["remove_chunks"]:
-                        document_chunks = remove_chunks(document_chunks)
-
-                    print(f"\t\tExtracted no. of chunks: {len(document_chunks)}")
-                    self.document_names.append(title)
-                    self.document_chunks_full.extend(document_chunks)
-                except:
-                    logger.info(f"\t\tError splitting link {link_index+1} : {link}")
-                    exit()
-
-        logger.info(
-            f"\tNumber of document chunks extracted in total: {len(self.document_chunks_full)}\n\n"
-        )
-
-        return self.document_chunks_full, self.document_names

code/modules/dataloader/data_loader.py

@@ -0,0 +1,360 @@
+import os
+import re
+import requests
+import pysrt
+from langchain_community.document_loaders import (
+    PyMuPDFLoader,
+    Docx2txtLoader,
+    YoutubeLoader,
+    WebBaseLoader,
+    TextLoader,
+)
+from langchain_community.document_loaders import UnstructuredMarkdownLoader
+from llama_parse import LlamaParse
+from langchain.schema import Document
+import logging
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from ragatouille import RAGPretrainedModel
+from langchain.chains import LLMChain
+from langchain_community.llms import OpenAI
+from langchain import PromptTemplate
+import json
+from concurrent.futures import ThreadPoolExecutor
+
+from modules.dataloader.helpers import get_metadata
+
+
+class PDFReader:
+    def __init__(self):
+        pass
+
+    def get_loader(self, pdf_path):
+        loader = PyMuPDFLoader(pdf_path)
+        return loader
+
+    def get_documents(self, loader):
+        return loader.load()
+
+
+class FileReader:
+    def __init__(self, logger):
+        self.pdf_reader = PDFReader()
+        self.logger = logger
+
+    def extract_text_from_pdf(self, pdf_path):
+        text = ""
+        with open(pdf_path, "rb") as file:
+            reader = PyPDF2.PdfReader(file)
+            num_pages = len(reader.pages)
+            for page_num in range(num_pages):
+                page = reader.pages[page_num]
+                text += page.extract_text()
+        return text
+
+    def download_pdf_from_url(self, pdf_url):
+        response = requests.get(pdf_url)
+        if response.status_code == 200:
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
+                temp_file.write(response.content)
+                temp_file_path = temp_file.name
+            return temp_file_path
+        else:
+            self.logger.error(f"Failed to download PDF from URL: {pdf_url}")
+            return None
+
+    def read_pdf(self, temp_file_path: str):
+        loader = self.pdf_reader.get_loader(temp_file_path)
+        documents = self.pdf_reader.get_documents(loader)
+        return documents
+
+    def read_txt(self, temp_file_path: str):
+        loader = TextLoader(temp_file_path, autodetect_encoding=True)
+        return loader.load()
+
+    def read_docx(self, temp_file_path: str):
+        loader = Docx2txtLoader(temp_file_path)
+        return loader.load()
+
+    def read_srt(self, temp_file_path: str):
+        subs = pysrt.open(temp_file_path)
+        text = ""
+        for sub in subs:
+            text += sub.text
+        return [Document(page_content=text)]
+
+    def read_youtube_transcript(self, url: str):
+        loader = YoutubeLoader.from_youtube_url(
+            url, add_video_info=True, language=["en"], translation="en"
+        )
+        return loader.load()
+
+    def read_html(self, url: str):
+        loader = WebBaseLoader(url)
+        return loader.load()
+
+    def read_tex_from_url(self, tex_url):
+        response = requests.get(tex_url)
+        if response.status_code == 200:
+            return [Document(page_content=response.text)]
+        else:
+            self.logger.error(f"Failed to fetch .tex file from URL: {tex_url}")
+            return None
+
+
+class ChunkProcessor:
+    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"]:
+                self.splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+                    chunk_size=config["splitter_options"]["chunk_size"],
+                    chunk_overlap=config["splitter_options"]["chunk_overlap"],
+                    separators=config["splitter_options"]["chunk_separators"],
+                    disallowed_special=(),
+                )
+            else:
+                self.splitter = RecursiveCharacterTextSplitter(
+                    chunk_size=config["splitter_options"]["chunk_size"],
+                    chunk_overlap=config["splitter_options"]["chunk_overlap"],
+                    separators=config["splitter_options"]["chunk_separators"],
+                    disallowed_special=(),
+                )
+        else:
+            self.splitter = None
+        self.logger.info("ChunkProcessor instance created")
+
+    def remove_delimiters(self, document_chunks: list):
+        for chunk in document_chunks:
+            for delimiter in self.config["splitter_options"]["delimiters_to_remove"]:
+                chunk.page_content = re.sub(delimiter, " ", chunk.page_content)
+        return document_chunks
+
+    def remove_chunks(self, document_chunks: list):
+        front = self.config["splitter_options"]["front_chunk_to_remove"]
+        end = self.config["splitter_options"]["last_chunks_to_remove"]
+        for _ in range(front):
+            del document_chunks[0]
+        for _ in range(end):
+            document_chunks.pop()
+        return document_chunks
+
+    def process_chunks(
+        self, documents, file_type="txt", source="", page=0, metadata={}
+    ):
+        documents = [Document(page_content=documents, source=source, page=page)]
+        if (
+            file_type == "txt"
+            or file_type == "docx"
+            or file_type == "srt"
+            or file_type == "tex"
+        ):
+            document_chunks = self.splitter.split_documents(documents)
+        elif file_type == "pdf":
+            document_chunks = documents  # Full page for now
+
+        # add the source and page number back to the metadata
+        for chunk in document_chunks:
+            chunk.metadata["source"] = source
+            chunk.metadata["page"] = page
+
+            # add the metadata extracted from the document
+            for key, value in metadata.items():
+                chunk.metadata[key] = value
+
+        if self.config["splitter_options"]["remove_leftover_delimiters"]:
+            document_chunks = self.remove_delimiters(document_chunks)
+        if self.config["splitter_options"]["remove_chunks"]:
+            document_chunks = self.remove_chunks(document_chunks)
+
+        return document_chunks
+
+    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
+
+        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, 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: # better approach to filter out non-informative documents
+            #     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, 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.chunk_docs(
+            self.file_reader, uploaded_files, weblinks
+        )
+
+
+if __name__ == "__main__":
+    import yaml
+
+    logger = logging.getLogger(__name__)
+    logger.setLevel(logging.INFO)
+
+    with open("../code/modules/config/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/"],
+        )
+    )
+
+    print(document_names)
+    print(len(document_chunks))

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/helpers.py

@@ -1,200 +0,0 @@
-import requests
-from bs4 import BeautifulSoup
-from tqdm import tqdm
-from urllib.parse import urlparse
-import chainlit as cl
-from langchain import PromptTemplate
-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):
-        pass
-
-    def getdata(self, url):
-        r = requests.get(url)
-        return r.text
-
-    def url_exists(self, url):
-        try:
-            response = requests.head(url)
-            return response.status_code == 200
-        except requests.ConnectionError:
-            return False
-
-    def get_links(self, website_link, base_url=None):
-        if base_url is None:
-            base_url = website_link
-        html_data = self.getdata(website_link)
-        soup = BeautifulSoup(html_data, "html.parser")
-        list_links = []
-        for link in soup.find_all("a", href=True):
-
-            # clean the link
-            # remove empty spaces
-            link["href"] = link["href"].strip()
-            # Append to list if new link contains original link
-            if str(link["href"]).startswith((str(website_link))):
-                list_links.append(link["href"])
-
-            # Include all href that do not start with website link but with "/"
-            if str(link["href"]).startswith("/"):
-                if link["href"] not in self.dict_href_links:
-                    print(link["href"])
-                    self.dict_href_links[link["href"]] = None
-                    link_with_www = base_url + link["href"][1:]
-                    if self.url_exists(link_with_www):
-                        print("adjusted link =", link_with_www)
-                        list_links.append(link_with_www)
-
-        # Convert list of links to dictionary and define keys as the links and the values as "Not-checked"
-        dict_links = dict.fromkeys(list_links, "Not-checked")
-        return dict_links
-
-    def get_subpage_links(self, l, base_url):
-        for link in tqdm(l):
-            print('checking link:', link)
-            if not link.endswith("/"):
-                l[link] = "Checked"
-                dict_links_subpages = {}
-            else:
-                # If not crawled through this page start crawling and get links
-                if l[link] == "Not-checked":
-                    dict_links_subpages = self.get_links(link, base_url)
-                    # Change the dictionary value of the link to "Checked"
-                    l[link] = "Checked"
-                else:
-                    # Create an empty dictionary in case every link is checked
-                    dict_links_subpages = {}
-            # Add new dictionary to old dictionary
-            l = {**dict_links_subpages, **l}
-        return l
-
-    def get_all_pages(self, url, base_url):
-        dict_links = {url: "Not-checked"}
-        self.dict_href_links = {}
-        counter, counter2 = None, 0
-        while counter != 0:
-            counter2 += 1
-            dict_links2 = self.get_subpage_links(dict_links, base_url)
-            # Count number of non-values and set counter to 0 if there are no values within the dictionary equal to the string "Not-checked"
-            # https://stackoverflow.com/questions/48371856/count-the-number-of-occurrences-of-a-certain-value-in-a-dictionary-in-python
-            counter = sum(value == "Not-checked" for value in dict_links2.values())
-            dict_links = dict_links2
-        checked_urls = [
-            url for url, status in dict_links.items() if status == "Checked"
-        ]
-        return checked_urls
-
-
-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
-
-def get_sources(res, answer):
-    source_elements_dict = {}
-    source_elements = []
-    found_sources = []
-
-    source_dict = {}  # Dictionary to store URL elements
-
-    for idx, source in enumerate(res["source_documents"]):
-        source_metadata = source.metadata
-        url = source_metadata["source"]
-
-        if url not in source_dict:
-            source_dict[url] = [source.page_content]
-        else:
-            source_dict[url].append(source.page_content)
-
-    for source_idx, (url, text_list) in enumerate(source_dict.items()):
-        full_text = ""
-        for url_idx, text in enumerate(text_list):
-            full_text += f"Source {url_idx+1}:\n {text}\n\n\n"
-        source_elements.append(cl.Text(name=url, content=full_text))
-        found_sources.append(url)
-
-    if found_sources:
-        answer += f"\n\nSources: {', '.join(found_sources)} "
-    else:
-        answer += f"\n\nNo source found."
-
-    # for idx, source in enumerate(res["source_documents"]):
-    #     title = source.metadata["source"]
-
-    #     if title not in source_elements_dict:
-    #         source_elements_dict[title] = {
-    #             "page_number": [source.metadata["page"]],
-    #             "url": source.metadata["source"],
-    #             "content": source.page_content,
-    #         }
-
-    #     else:
-    #         source_elements_dict[title]["page_number"].append(source.metadata["page"])
-    #     source_elements_dict[title][
-    #         "content_" + str(source.metadata["page"])
-    #     ] = source.page_content
-    #     # sort the page numbers
-    #     # source_elements_dict[title]["page_number"].sort()
-
-    # for title, source in source_elements_dict.items():
-    #     # create a string for the page numbers
-    #     page_numbers = ", ".join([str(x) for x in source["page_number"]])
-    #     text_for_source = f"Page Number(s): {page_numbers}\nURL: {source['url']}"
-    #     source_elements.append(cl.Pdf(name="File", path=title))
-    #     found_sources.append("File")
-    #     # for pn in source["page_number"]:
-    #     #     source_elements.append(
-    #     #         cl.Text(name=str(pn), content=source["content_"+str(pn)])
-    #     #     )
-    #     #     found_sources.append(str(pn))
-
-    # if found_sources:
-    #     answer += f"\nSource:{', '.join(found_sources)}"
-    # else:
-    #     answer += f"\nNo source found."
-
-    return answer, source_elements

code/modules/llm_tutor.py

@@ -1,87 +0,0 @@
-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
-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
-
-
-class LLMTutor:
-    def __init__(self, config, logger=None):
-        self.config = config
-        self.vector_db = VectorDB(config, logger=logger)
-        if self.config["embedding_options"]["embedd_files"]:
-            self.vector_db.create_database()
-            self.vector_db.save_database()
-
-    def set_custom_prompt(self):
-        """
-        Prompt template for QA retrieval for each vectorstore
-        """
-        prompt = get_prompt(self.config)
-        # prompt = QA_PROMPT
-
-        return prompt
-
-    # Retrieval QA Chain
-    def retrieval_qa_chain(self, llm, prompt, db):
-        if self.config["llm_params"]["use_history"]:
-            memory = ConversationBufferWindowMemory(
-            k = self.config["llm_params"]["memory_window"], 
-            memory_key="chat_history", return_messages=True, output_key="answer"
-            )
-            qa_chain = ConversationalRetrievalChain.from_llm(
-                llm=llm,
-                chain_type="stuff",
-                retriever=db.as_retriever(
-                    search_kwargs={
-                        "k": self.config["embedding_options"]["search_top_k"]
-                    }
-                ),
-                return_source_documents=True,
-                memory=memory,
-                combine_docs_chain_kwargs={"prompt": prompt},
-            )
-        else:
-            qa_chain = RetrievalQA.from_chain_type(
-                llm=llm,
-                chain_type="stuff",
-                retriever=db.as_retriever(
-                    search_kwargs={
-                        "k": self.config["embedding_options"]["search_top_k"]
-                    }
-                ),
-                return_source_documents=True,
-                chain_type_kwargs={"prompt": prompt},
-            )
-        return qa_chain
-
-    # Loading the model
-    def load_llm(self):
-        chat_model_loader = ChatModelLoader(self.config)
-        llm = chat_model_loader.load_chat_model()
-        return llm
-
-    # QA Model Function
-    def qa_bot(self):
-        db = self.vector_db.load_database()
-        self.llm = self.load_llm()
-        qa_prompt = self.set_custom_prompt()
-        qa = self.retrieval_qa_chain(self.llm, qa_prompt, db)
-
-        return qa
-
-    # output function
-    def final_result(query):
-        qa_result = qa_bot()
-        response = qa_result({"query": query})
-        return response

code/modules/retriever/base.py

@@ -0,0 +1,12 @@
+# template for retriever classes
+
+
+class BaseRetriever:
+    def __init__(self, config):
+        self.config = config
+
+    def return_retriever(self):
+        """
+        Returns the retriever object
+        """
+        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/colbert_retriever.py

@@ -0,0 +1,10 @@
+from .base import BaseRetriever
+
+
+class ColbertRetriever(BaseRetriever):
+    def __init__(self):
+        pass
+
+    def return_retriever(self, db, config):
+        retriever = db.as_langchain_retriever(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/retriever/raptor_retriever.py

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

code/modules/retriever/retriever.py

@@ -0,0 +1,26 @@
+from modules.retriever.faiss_retriever import FaissRetriever
+from modules.retriever.chroma_retriever import ChromaRetriever
+from modules.retriever.colbert_retriever import ColbertRetriever
+from modules.retriever.raptor_retriever import RaptorRetriever
+
+
+class Retriever:
+    def __init__(self, config):
+        self.config = config
+        self.retriever_classes = {
+            "FAISS": FaissRetriever,
+            "Chroma": ChromaRetriever,
+            "RAGatouille": ColbertRetriever,
+            "RAPTOR": RaptorRetriever,
+        }
+        self._create_retriever()
+
+    def _create_retriever(self):
+        db_option = self.config["vectorstore"]["db_option"]
+        retriever_class = self.retriever_classes.get(db_option)
+        if not retriever_class:
+            raise ValueError(f"Invalid db_option: {db_option}")
+        self.retriever = retriever_class()
+
+    def _return_retriever(self, db):
+        return self.retriever.return_retriever(db, self.config)

code/modules/vector_db.py

@@ -1,133 +0,0 @@
-import logging
-import os
-import yaml
-from langchain.vectorstores import FAISS
-
-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 *
-
-
-class VectorDB:
-    def __init__(self, config, logger=None):
-        self.config = config
-        self.db_option = config["embedding_options"]["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)
-
-            # Console Handler
-            console_handler = logging.StreamHandler()
-            console_handler.setLevel(logging.INFO)
-            formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
-            console_handler.setFormatter(formatter)
-            self.logger.addHandler(console_handler)
-
-            # File Handler
-            log_file_path = "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)
-            self.logger.addHandler(file_handler)
-        else:
-            self.logger = logger
-
-        self.logger.info("VectorDB instance instantiated")
-
-    def load_files(self):
-        files = os.listdir(self.config["embedding_options"]["data_path"])
-        files = [
-            os.path.join(self.config["embedding_options"]["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"]:
-            all_urls = []
-            for url in urls:
-                base_url = get_base_url(url)
-                all_urls.extend(self.webpage_crawler.get_all_pages(url, base_url))
-            urls = all_urls
-        return files, urls
-
-    def clean_url_list(self, urls):
-        # get lecture pdf links 
-        lecture_pdfs = [link for link in urls if link.endswith(".pdf")]
-        lecture_pdfs = [link for link in lecture_pdfs if "lecture" in link.lower()]
-        urls = [link for link in urls if link.endswith("/")] # only keep links that end with a '/'. Extract Files Seperately
-
-        return urls, lecture_pdfs
-
-    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()
-
-    def initialize_database(self, document_chunks: list, document_names: list):
-        # Track token usage
-        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.logger.info("Completed initializing vector_db")
-
-    def create_database(self):
-        data_loader = DataLoader(self.config)
-        self.logger.info("Loading data")
-        files, urls = self.load_files()
-        urls, lecture_pdfs = self.clean_url_list(urls)
-        files += lecture_pdfs
-        files.remove('storage/data/urls.txt')
-        document_chunks, document_names = data_loader.get_chunks(files, urls)
-        self.logger.info("Completed loading data")
-
-        self.create_embedding_model()
-        self.initialize_database(document_chunks, document_names)
-
-    def save_database(self):
-        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"],
-            )
-        )
-        self.logger.info("Saved database")
-
-    def load_database(self):
-        self.create_embedding_model()
-        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,
-        )
-        self.logger.info("Loaded database")
-        return self.vector_db
-
-
-if __name__ == "__main__":
-    with open("code/config.yml", "r") as f:
-        config = yaml.safe_load(f)
-    print(config)
-    vector_db = VectorDB(config)
-    vector_db.create_database()
-    vector_db.save_database()

code/modules/vectorstore/base.py

@@ -0,0 +1,33 @@
+# template for vector store classes
+
+
+class VectorStoreBase:
+    def __init__(self, config):
+        self.config = config
+
+    def _init_vector_db(self):
+        """
+        Creates a vector store object
+        """
+        raise NotImplementedError
+
+    def create_database(self):
+        """
+        Populates the vector store with documents
+        """
+        raise NotImplementedError
+
+    def load_database(self):
+        """
+        Loads the vector store from disk
+        """
+        raise NotImplementedError
+
+    def as_retriever(self):
+        """
+        Returns the vector store as a retriever
+        """
+        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/vectorstore/colbert.py

@@ -0,0 +1,39 @@
+from ragatouille import RAGPretrainedModel
+from modules.vectorstore.base import VectorStoreBase
+import os
+
+
+class ColbertVectorStore(VectorStoreBase):
+    def __init__(self, config):
+        self.config = config
+        self._init_vector_db()
+
+    def _init_vector_db(self):
+        self.colbert = RAGPretrainedModel.from_pretrained(
+            "colbert-ir/colbertv2.0",
+            index_root=os.path.join(
+                self.config["vectorstore"]["db_path"],
+                "db_" + self.config["vectorstore"]["db_option"],
+            ),
+        )
+
+    def create_database(self, documents, document_names, document_metadata):
+        index_path = self.colbert.index(
+            index_name="new_idx",
+            collection=documents,
+            document_ids=document_names,
+            document_metadatas=document_metadata,
+        )
+
+    def load_database(self):
+        path = os.path.join(
+            self.config["vectorstore"]["db_path"],
+            "db_" + self.config["vectorstore"]["db_option"],
+        )
+        self.vectorstore = RAGPretrainedModel.from_index(
+            f"{path}/colbert/indexes/new_idx"
+        )
+        return self.vectorstore
+
+    def as_retriever(self):
+        return self.vectorstore.as_retriever()

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

@@ -1,10 +1,8 @@
 from langchain_community.embeddings import OpenAIEmbeddings
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.embeddings import LlamaCppEmbeddings
-try:
-    from modules.constants import *
-except:
-    from constants import *
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.embeddings import LlamaCppEmbeddings
+
+from modules.config.constants import *
 import os
 
 
@@ -13,17 +11,22 @@ 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="sentence-transformers/all-MiniLM-L6-v2",
-                model_kwargs={"device": "cpu"},
+                model_name=self.config["vectorstore"]["model"],
+                model_kwargs={
+                    "device": f"{self.config['device']}",
+                    "token": f"{HUGGINGFACE_TOKEN}",
+                    "trust_remote_code": True,
+                },
             )
             # embedding_model = LlamaCppEmbeddings(
             #     model_path=os.path.abspath("storage/llama-7b.ggmlv3.q4_0.bin")

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/vectorstore/raptor.py

@@ -0,0 +1,438 @@
+# code modified from https://github.com/langchain-ai/langchain/blob/master/cookbook/RAPTOR.ipynb
+
+from typing import Dict, List, Optional, Tuple
+import os
+import numpy as np
+import pandas as pd
+import umap
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from sklearn.mixture import GaussianMixture
+from langchain_community.chat_models import ChatOpenAI
+from langchain_community.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from modules.vectorstore.base import VectorStoreBase
+
+RANDOM_SEED = 42
+
+
+class RAPTORVectoreStore(VectorStoreBase):
+    def __init__(self, config, documents=[], text_splitter=None, embedding_model=None):
+        self.documents = documents
+        self.config = config
+        self.text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+            chunk_size=self.config["splitter_options"]["chunk_size"],
+            chunk_overlap=self.config["splitter_options"]["chunk_overlap"],
+            separators=self.config["splitter_options"]["chunk_separators"],
+            disallowed_special=(),
+        )
+        self.embd = embedding_model
+        self.model = ChatOpenAI(
+            model="gpt-3.5-turbo",
+        )
+
+    def concat_documents(self, documents):
+        d_sorted = sorted(documents, key=lambda x: x.metadata["source"])
+        d_reversed = list(reversed(d_sorted))
+        concatenated_content = "\n\n\n --- \n\n\n".join(
+            [doc.page_content for doc in d_reversed]
+        )
+        return concatenated_content
+
+    def split_documents(self, documents):
+        concatenated_content = self.concat_documents(documents)
+        texts_split = self.text_splitter.split_text(concatenated_content)
+        return texts_split
+
+    def add_documents(self, documents):
+        self.documents.extend(documents)
+
+    def global_cluster_embeddings(
+        self,
+        embeddings: np.ndarray,
+        dim: int,
+        n_neighbors: Optional[int] = None,
+        metric: str = "cosine",
+    ) -> np.ndarray:
+        """
+        Perform global dimensionality reduction on the embeddings using UMAP.
+
+        Parameters:
+        - embeddings: The input embeddings as a numpy array.
+        - dim: The target dimensionality for the reduced space.
+        - n_neighbors: Optional; the number of neighbors to consider for each point.
+                    If not provided, it defaults to the square root of the number of embeddings.
+        - metric: The distance metric to use for UMAP.
+
+        Returns:
+        - A numpy array of the embeddings reduced to the specified dimensionality.
+        """
+        if n_neighbors is None:
+            n_neighbors = int((len(embeddings) - 1) ** 0.5)
+        return umap.UMAP(
+            n_neighbors=n_neighbors, n_components=dim, metric=metric
+        ).fit_transform(embeddings)
+
+    def local_cluster_embeddings(
+        self,
+        embeddings: np.ndarray,
+        dim: int,
+        num_neighbors: int = 10,
+        metric: str = "cosine",
+    ) -> np.ndarray:
+        """
+        Perform local dimensionality reduction on the embeddings using UMAP, typically after global clustering.
+
+        Parameters:
+        - embeddings: The input embeddings as a numpy array.
+        - dim: The target dimensionality for the reduced space.
+        - num_neighbors: The number of neighbors to consider for each point.
+        - metric: The distance metric to use for UMAP.
+
+        Returns:
+        - A numpy array of the embeddings reduced to the specified dimensionality.
+        """
+        return umap.UMAP(
+            n_neighbors=num_neighbors, n_components=dim, metric=metric
+        ).fit_transform(embeddings)
+
+    def get_optimal_clusters(
+        self,
+        embeddings: np.ndarray,
+        max_clusters: int = 50,
+        random_state: int = RANDOM_SEED,
+    ) -> int:
+        """
+        Determine the optimal number of clusters using the Bayesian Information Criterion (BIC) with a Gaussian Mixture Model.
+
+        Parameters:
+        - embeddings: The input embeddings as a numpy array.
+        - max_clusters: The maximum number of clusters to consider.
+        - random_state: Seed for reproducibility.
+
+        Returns:
+        - An integer representing the optimal number of clusters found.
+        """
+        max_clusters = min(max_clusters, len(embeddings))
+        n_clusters = np.arange(1, max_clusters)
+        bics = []
+        for n in n_clusters:
+            gm = GaussianMixture(n_components=n, random_state=random_state)
+            gm.fit(embeddings)
+            bics.append(gm.bic(embeddings))
+        return n_clusters[np.argmin(bics)]
+
+    def GMM_cluster(
+        self, embeddings: np.ndarray, threshold: float, random_state: int = 0
+    ):
+        """
+        Cluster embeddings using a Gaussian Mixture Model (GMM) based on a probability threshold.
+
+        Parameters:
+        - embeddings: The input embeddings as a numpy array.
+        - threshold: The probability threshold for assigning an embedding to a cluster.
+        - random_state: Seed for reproducibility.
+
+        Returns:
+        - A tuple containing the cluster labels and the number of clusters determined.
+        """
+        n_clusters = self.get_optimal_clusters(embeddings)
+        gm = GaussianMixture(n_components=n_clusters, random_state=random_state)
+        gm.fit(embeddings)
+        probs = gm.predict_proba(embeddings)
+        labels = [np.where(prob > threshold)[0] for prob in probs]
+        return labels, n_clusters
+
+    def perform_clustering(
+        self,
+        embeddings: np.ndarray,
+        dim: int,
+        threshold: float,
+    ) -> List[np.ndarray]:
+        """
+        Perform clustering on the embeddings by first reducing their dimensionality globally, then clustering
+        using a Gaussian Mixture Model, and finally performing local clustering within each global cluster.
+
+        Parameters:
+        - embeddings: The input embeddings as a numpy array.
+        - dim: The target dimensionality for UMAP reduction.
+        - threshold: The probability threshold for assigning an embedding to a cluster in GMM.
+
+        Returns:
+        - A list of numpy arrays, where each array contains the cluster IDs for each embedding.
+        """
+        if len(embeddings) <= dim + 1:
+            # Avoid clustering when there's insufficient data
+            return [np.array([0]) for _ in range(len(embeddings))]
+
+        # Global dimensionality reduction
+        reduced_embeddings_global = self.global_cluster_embeddings(embeddings, dim)
+        # Global clustering
+        global_clusters, n_global_clusters = self.GMM_cluster(
+            reduced_embeddings_global, threshold
+        )
+
+        all_local_clusters = [np.array([]) for _ in range(len(embeddings))]
+        total_clusters = 0
+
+        # Iterate through each global cluster to perform local clustering
+        for i in range(n_global_clusters):
+            # Extract embeddings belonging to the current global cluster
+            global_cluster_embeddings_ = embeddings[
+                np.array([i in gc for gc in global_clusters])
+            ]
+
+            if len(global_cluster_embeddings_) == 0:
+                continue
+            if len(global_cluster_embeddings_) <= dim + 1:
+                # Handle small clusters with direct assignment
+                local_clusters = [np.array([0]) for _ in global_cluster_embeddings_]
+                n_local_clusters = 1
+            else:
+                # Local dimensionality reduction and clustering
+                reduced_embeddings_local = self.local_cluster_embeddings(
+                    global_cluster_embeddings_, dim
+                )
+                local_clusters, n_local_clusters = self.GMM_cluster(
+                    reduced_embeddings_local, threshold
+                )
+
+            # Assign local cluster IDs, adjusting for total clusters already processed
+            for j in range(n_local_clusters):
+                local_cluster_embeddings_ = global_cluster_embeddings_[
+                    np.array([j in lc for lc in local_clusters])
+                ]
+                indices = np.where(
+                    (embeddings == local_cluster_embeddings_[:, None]).all(-1)
+                )[1]
+                for idx in indices:
+                    all_local_clusters[idx] = np.append(
+                        all_local_clusters[idx], j + total_clusters
+                    )
+
+            total_clusters += n_local_clusters
+
+        return all_local_clusters
+
+    def embed(self, texts):
+        """
+        Generate embeddings for a list of text documents.
+
+        This function assumes the existence of an `embd` object with a method `embed_documents`
+        that takes a list of texts and returns their embeddings.
+
+        Parameters:
+        - texts: List[str], a list of text documents to be embedded.
+
+        Returns:
+        - numpy.ndarray: An array of embeddings for the given text documents.
+        """
+        text_embeddings = self.embd.embed_documents(texts)
+        text_embeddings_np = np.array(text_embeddings)
+        return text_embeddings_np
+
+    def embed_cluster_texts(self, texts):
+        """
+        Embeds a list of texts and clusters them, returning a DataFrame with texts, their embeddings, and cluster labels.
+
+        This function combines embedding generation and clustering into a single step. It assumes the existence
+        of a previously defined `perform_clustering` function that performs clustering on the embeddings.
+
+        Parameters:
+        - texts: List[str], a list of text documents to be processed.
+
+        Returns:
+        - pandas.DataFrame: A DataFrame containing the original texts, their embeddings, and the assigned cluster labels.
+        """
+        text_embeddings_np = self.embed(texts)  # Generate embeddings
+        cluster_labels = self.perform_clustering(
+            text_embeddings_np, 10, 0.1
+        )  # Perform clustering on the embeddings
+        df = pd.DataFrame()  # Initialize a DataFrame to store the results
+        df["text"] = texts  # Store original texts
+        df["embd"] = list(
+            text_embeddings_np
+        )  # Store embeddings as a list in the DataFrame
+        df["cluster"] = cluster_labels  # Store cluster labels
+        return df
+
+    def fmt_txt(self, df: pd.DataFrame) -> str:
+        """
+        Formats the text documents in a DataFrame into a single string.
+
+        Parameters:
+        - df: DataFrame containing the 'text' column with text documents to format.
+
+        Returns:
+        - A single string where all text documents are joined by a specific delimiter.
+        """
+        unique_txt = df["text"].tolist()
+        return "--- --- \n --- --- ".join(unique_txt)
+
+    def embed_cluster_summarize_texts(
+        self, texts: List[str], level: int
+    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
+        """
+        Embeds, clusters, and summarizes a list of texts. This function first generates embeddings for the texts,
+        clusters them based on similarity, expands the cluster assignments for easier processing, and then summarizes
+        the content within each cluster.
+
+        Parameters:
+        - texts: A list of text documents to be processed.
+        - level: An integer parameter that could define the depth or detail of processing.
+
+        Returns:
+        - Tuple containing two DataFrames:
+        1. The first DataFrame (`df_clusters`) includes the original texts, their embeddings, and cluster assignments.
+        2. The second DataFrame (`df_summary`) contains summaries for each cluster, the specified level of detail,
+            and the cluster identifiers.
+        """
+
+        # Embed and cluster the texts, resulting in a DataFrame with 'text', 'embd', and 'cluster' columns
+        df_clusters = self.embed_cluster_texts(texts)
+
+        # Prepare to expand the DataFrame for easier manipulation of clusters
+        expanded_list = []
+
+        # Expand DataFrame entries to document-cluster pairings for straightforward processing
+        for index, row in df_clusters.iterrows():
+            for cluster in row["cluster"]:
+                expanded_list.append(
+                    {"text": row["text"], "embd": row["embd"], "cluster": cluster}
+                )
+
+        # Create a new DataFrame from the expanded list
+        expanded_df = pd.DataFrame(expanded_list)
+
+        # Retrieve unique cluster identifiers for processing
+        all_clusters = expanded_df["cluster"].unique()
+
+        print(f"--Generated {len(all_clusters)} clusters--")
+
+        # Summarization
+        template = """Here is content from the course DS598: Deep Learning for Data Science. 
+        
+        The content may be form webapge about the course, or lecture content, or any other relevant information.
+        If the content is in bullet points (from  pdf lectre slides), you can summarize the bullet points.
+        
+        Give a detailed summary of the content below.
+        
+        Documentation:
+        {context}
+        """
+        prompt = ChatPromptTemplate.from_template(template)
+        chain = prompt | self.model | StrOutputParser()
+
+        # Format text within each cluster for summarization
+        summaries = []
+        for i in all_clusters:
+            df_cluster = expanded_df[expanded_df["cluster"] == i]
+            formatted_txt = self.fmt_txt(df_cluster)
+            summaries.append(chain.invoke({"context": formatted_txt}))
+
+        # Create a DataFrame to store summaries with their corresponding cluster and level
+        df_summary = pd.DataFrame(
+            {
+                "summaries": summaries,
+                "level": [level] * len(summaries),
+                "cluster": list(all_clusters),
+            }
+        )
+
+        return df_clusters, df_summary
+
+    def recursive_embed_cluster_summarize(
+        self, texts: List[str], level: int = 1, n_levels: int = 3
+    ) -> Dict[int, Tuple[pd.DataFrame, pd.DataFrame]]:
+        """
+        Recursively embeds, clusters, and summarizes texts up to a specified level or until
+        the number of unique clusters becomes 1, storing the results at each level.
+
+        Parameters:
+        - texts: List[str], texts to be processed.
+        - level: int, current recursion level (starts at 1).
+        - n_levels: int, maximum depth of recursion.
+
+        Returns:
+        - Dict[int, Tuple[pd.DataFrame, pd.DataFrame]], a dictionary where keys are the recursion
+        levels and values are tuples containing the clusters DataFrame and summaries DataFrame at that level.
+        """
+        results = {}  # Dictionary to store results at each level
+
+        # Perform embedding, clustering, and summarization for the current level
+        df_clusters, df_summary = self.embed_cluster_summarize_texts(texts, level)
+
+        # Store the results of the current level
+        results[level] = (df_clusters, df_summary)
+
+        # Determine if further recursion is possible and meaningful
+        unique_clusters = df_summary["cluster"].nunique()
+        if level < n_levels and unique_clusters > 1:
+            # Use summaries as the input texts for the next level of recursion
+            new_texts = df_summary["summaries"].tolist()
+            next_level_results = self.recursive_embed_cluster_summarize(
+                new_texts, level + 1, n_levels
+            )
+
+            # Merge the results from the next level into the current results dictionary
+            results.update(next_level_results)
+
+        return results
+
+    def get_vector_db(self):
+        """
+        Generate a retriever object from a list of documents.
+
+        Parameters:
+        - documents: List of document objects.
+
+        Returns:
+        - A retriever object.
+        """
+        leaf_texts = self.split_documents(self.documents)
+        results = self.recursive_embed_cluster_summarize(
+            leaf_texts, level=1, n_levels=10
+        )
+
+        all_texts = leaf_texts.copy()
+        # Iterate through the results to extract summaries from each level and add them to all_texts
+        for level in sorted(results.keys()):
+            # Extract summaries from the current level's DataFrame
+            summaries = results[level][1]["summaries"].tolist()
+            # Extend all_texts with the summaries from the current level
+            all_texts.extend(summaries)
+
+        # Now, use all_texts to build the vectorstore
+        vectorstore = FAISS.from_texts(texts=all_texts, embedding=self.embd)
+        return vectorstore
+
+    def create_database(self, documents, embedding_model):
+        self.documents = documents
+        self.embd = embedding_model
+        self.vectorstore = self.get_vector_db()
+        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 = 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/vectorstore/store_manager.py

@@ -0,0 +1,163 @@
+from modules.vectorstore.vectorstore import VectorStore
+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 time
+import asyncio
+
+
+class VectorStoreManager:
+    def __init__(self, config, logger=None):
+        self.config = config
+        self.document_names = None
+
+        # Set up logging to both console and a file
+        self.logger = logger or self._setup_logging()
+        self.webpage_crawler = WebpageCrawler()
+        self.vector_db = VectorStore(self.config)
+
+        self.logger.info("VectorDB instance instantiated")
+
+    def _setup_logging(self):
+        logger = logging.getLogger(__name__)
+        if not logger.hasHandlers():
+            logger.setLevel(logging.INFO)
+            formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+
+            # Console Handler
+            console_handler = logging.StreamHandler()
+            console_handler.setLevel(logging.INFO)
+            console_handler.setFormatter(formatter)
+            logger.addHandler(console_handler)
+
+            # Ensure log directory exists
+            log_directory = self.config["log_dir"]
+            os.makedirs(log_directory, exist_ok=True)
+
+            # File Handler
+            log_file_path = os.path.join(log_directory, "vector_db.log")
+            file_handler = logging.FileHandler(log_file_path, mode="w")
+            file_handler.setLevel(logging.INFO)
+            file_handler.setFormatter(formatter)
+            logger.addHandler(file_handler)
+
+        return logger
+
+    def load_files(self):
+
+        files = os.listdir(self.config["vectorstore"]["data_path"])
+        files = [
+            os.path.join(self.config["vectorstore"]["data_path"], file)
+            for file in files
+        ]
+        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()
+                all_urls.extend(
+                    loop.run_until_complete(
+                        self.webpage_crawler.get_all_pages(
+                            url, url
+                        )  # only get child urls, if you want to get all urls, replace the second argument with the base url
+                    )
+                )
+            urls = all_urls
+        return files, urls
+
+    def create_embedding_model(self):
+
+        self.logger.info("Creating embedding function")
+        embedding_model_loader = EmbeddingModelLoader(self.config)
+        embedding_model = embedding_model_loader.load_embedding_model()
+        return embedding_model
+
+    def initialize_database(
+        self,
+        document_chunks: list,
+        document_names: list,
+        documents: list,
+        document_metadata: list,
+    ):
+        if self.config["vectorstore"]["db_option"] in ["FAISS", "Chroma", "RAPTOR"]:
+            self.embedding_model = self.create_embedding_model()
+        else:
+            self.embedding_model = None
+
+        self.logger.info("Initializing vector_db")
+        self.logger.info(
+            "\tUsing {} as db_option".format(self.config["vectorstore"]["db_option"])
+        )
+        self.vector_db._create_database(
+            document_chunks,
+            document_names,
+            documents,
+            document_metadata,
+            self.embedding_model,
+        )
+
+    def create_database(self):
+
+        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)
+        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 load_database(self):
+
+        start_time = time.time()  # Start time for loading database
+        if self.config["vectorstore"]["db_option"] in ["FAISS", "Chroma", "RAPTOR"]:
+            self.embedding_model = self.create_embedding_model()
+        else:
+            self.embedding_model = None
+        self.loaded_vector_db = self.vector_db._load_database(self.embedding_model)
+        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.loaded_vector_db
+
+
+if __name__ == "__main__":
+    import yaml
+
+    with open("modules/config/config.yml", "r") as f:
+        config = yaml.safe_load(f)
+    print(config)
+    print(f"Trying to create database with config: {config}")
+    vector_db = VectorStoreManager(config)
+    vector_db.create_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")

code/modules/vectorstore/vectorstore.py

@@ -0,0 +1,57 @@
+from modules.vectorstore.faiss import FaissVectorStore
+from modules.vectorstore.chroma import ChromaVectorStore
+from modules.vectorstore.colbert import ColbertVectorStore
+from modules.vectorstore.raptor import RAPTORVectoreStore
+
+
+class VectorStore:
+    def __init__(self, config):
+        self.config = config
+        self.vectorstore = None
+        self.vectorstore_classes = {
+            "FAISS": FaissVectorStore,
+            "Chroma": ChromaVectorStore,
+            "RAGatouille": ColbertVectorStore,
+            "RAPTOR": RAPTORVectoreStore,
+        }
+
+    def _create_database(
+        self,
+        document_chunks,
+        document_names,
+        documents,
+        document_metadata,
+        embedding_model,
+    ):
+        db_option = self.config["vectorstore"]["db_option"]
+        vectorstore_class = self.vectorstore_classes.get(db_option)
+        if not vectorstore_class:
+            raise ValueError(f"Invalid db_option: {db_option}")
+
+        self.vectorstore = vectorstore_class(self.config)
+
+        if db_option == "RAGatouille":
+            self.vectorstore.create_database(
+                documents, document_names, document_metadata
+            )
+        else:
+            self.vectorstore.create_database(document_chunks, embedding_model)
+
+    def _load_database(self, embedding_model):
+        db_option = self.config["vectorstore"]["db_option"]
+        vectorstore_class = self.vectorstore_classes.get(db_option)
+        if not vectorstore_class:
+            raise ValueError(f"Invalid db_option: {db_option}")
+
+        self.vectorstore = vectorstore_class(self.config)
+
+        if db_option == "RAGatouille":
+            return self.vectorstore.load_database()
+        else:
+            return self.vectorstore.load_database(embedding_model)
+
+    def _as_retriever(self):
+        return self.vectorstore.as_retriever()
+
+    def _get_vectorstore(self):
+        return self.vectorstore