Merge pull request #36 from DL4DS/chainlit_enhancements

Dockerfile.dev

@@ -25,7 +25,7 @@ RUN mkdir /.cache && chmod -R 777 /.cache
 WORKDIR /code/code
 
 # Expose the port the app runs on
-EXPOSE 8051
+EXPOSE 8000
 
 # 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"]
+CMD ["sh", "-c", "python -m modules.vectorstore.store_manager && chainlit run main.py --host 0.0.0.0 --port 8000"]

README.md

@@ -76,7 +76,7 @@ The HuggingFace Space is built using the `Dockerfile` in the repository. To run
 
 ```bash
 docker build --tag dev  -f Dockerfile.dev .
-docker run -it --rm -p 8051:8051 dev
+docker run -it --rm -p 8000:8000 dev
 ```
 
 ## Contributing

code/main.py

@@ -1,178 +1,455 @@
-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
-import chainlit as cl
-from langchain_community.chat_models import ChatOpenAI
-from langchain_community.embeddings import OpenAIEmbeddings
-import yaml
-import logging
-from dotenv import load_dotenv
+import chainlit.data as cl_data
+import asyncio
+from modules.config.constants import (
+    LLAMA_PATH,
+    LITERAL_API_KEY_LOGGING,
+    LITERAL_API_URL,
+)
+from modules.chat_processor.literal_ai import CustomLiteralDataLayer
 
+import json
+import yaml
+import os
+from typing import Any, Dict, no_type_check
+import chainlit as cl
 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)
-console_handler.setFormatter(formatter)
-logger.addHandler(console_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="Mistral",
-        #     markdown_description="Use the local LLM: **Mistral**.",
-        # ),
-        cl.ChatProfile(
-            name="gpt-3.5-turbo-1106",
-            markdown_description="Use OpenAI API for **gpt-3.5-turbo-1106**.",
-        ),
-        cl.ChatProfile(
-            name="gpt-4",
-            markdown_description="Use OpenAI API for **gpt-4**.",
-        ),
-        cl.ChatProfile(
-            name="Llama",
-            markdown_description="Use the local LLM: **Tiny Llama**.",
-        ),
-    ]
-
-
-@cl.author_rename
-def rename(orig_author: str):
-    rename_dict = {"Chatbot": "AI Tutor"}
-    return rename_dict.get(orig_author, orig_author)
-
-
-# chainlit code
-@cl.on_chat_start
-async def start():
-    with open("modules/config/config.yml", "r") as f:
-        config = yaml.safe_load(f)
-
-    # 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:
-        if chat_profile.lower() in ["gpt-3.5-turbo-1106", "gpt-4"]:
-            config["llm_params"]["llm_loader"] = "openai"
-            config["llm_params"]["openai_params"]["model"] = chat_profile.lower()
-        elif chat_profile.lower() == "llama":
-            config["llm_params"]["llm_loader"] = "local_llm"
-            config["llm_params"]["local_llm_params"]["model"] = LLAMA_PATH
-            config["llm_params"]["local_llm_params"]["model_type"] = "llama"
-        elif chat_profile.lower() == "mistral":
-            config["llm_params"]["llm_loader"] = "local_llm"
-            config["llm_params"]["local_llm_params"]["model"] = MISTRAL_PATH
-            config["llm_params"]["local_llm_params"]["model_type"] = "mistral"
+from modules.chat.helpers import (
+    get_sources,
+    get_history_chat_resume,
+    get_history_setup_llm,
+    get_last_config,
+)
+import copy
+from typing import Optional
+from chainlit.types import ThreadDict
+import time
+
+USER_TIMEOUT = 60_000
+SYSTEM = "System 🖥️"
+LLM = "LLM 🧠"
+AGENT = "Agent <>"
+YOU = "You 😃"
+ERROR = "Error 🚫"
+
+with open("modules/config/config.yml", "r") as f:
+    config = yaml.safe_load(f)
+
+
+async def setup_data_layer():
+    """
+    Set up the data layer for chat logging.
+    """
+    if config["chat_logging"]["log_chat"]:
+        data_layer = CustomLiteralDataLayer(
+            api_key=LITERAL_API_KEY_LOGGING, server=LITERAL_API_URL
+        )
+    else:
+        data_layer = None
+
+    return data_layer
+
+
+class Chatbot:
+    def __init__(self, config):
+        """
+        Initialize the Chatbot class.
+        """
+        self.config = config
+
+    async def _load_config(self):
+        """
+        Load the configuration from a YAML file.
+        """
+        with open("modules/config/config.yml", "r") as f:
+            return yaml.safe_load(f)
+
+    @no_type_check
+    async def setup_llm(self):
+        """
+        Set up the LLM with the provided settings. Update the configuration and initialize the LLM tutor.
+        """
+        start_time = time.time()
+
+        llm_settings = cl.user_session.get("llm_settings", {})
+        chat_profile, retriever_method, memory_window, llm_style, generate_follow_up = (
+            llm_settings.get("chat_model"),
+            llm_settings.get("retriever_method"),
+            llm_settings.get("memory_window"),
+            llm_settings.get("llm_style"),
+            llm_settings.get("follow_up_questions"),
+        )
+
+        chain = cl.user_session.get("chain")
+        memory_list = cl.user_session.get(
+            "memory",
+            (
+                list(chain.store.values())[0].messages
+                if len(chain.store.values()) > 0
+                else []
+            ),
+        )
+        conversation_list = get_history_setup_llm(memory_list)
+
+        old_config = copy.deepcopy(self.config)
+        self.config["vectorstore"]["db_option"] = retriever_method
+        self.config["llm_params"]["memory_window"] = memory_window
+        self.config["llm_params"]["llm_style"] = llm_style
+        self.config["llm_params"]["llm_loader"] = chat_profile
+        self.config["llm_params"]["generate_follow_up"] = generate_follow_up
+
+        self.llm_tutor.update_llm(
+            old_config, self.config
+        )  # update only llm attributes that are changed
+        self.chain = self.llm_tutor.qa_bot(
+            memory=conversation_list,
+            callbacks=(
+                [cl.LangchainCallbackHandler()]
+                if cl_data._data_layer and self.config["chat_logging"]["callbacks"]
+                else None
+            ),
+        )
+
+        tags = [chat_profile, self.config["vectorstore"]["db_option"]]
+
+        cl.user_session.set("chain", self.chain)
+        cl.user_session.set("llm_tutor", self.llm_tutor)
+
+        print("Time taken to setup LLM: ", time.time() - start_time)
+
+    @no_type_check
+    async def update_llm(self, new_settings: Dict[str, Any]):
+        """
+        Update the LLM settings and reinitialize the LLM with the new settings.
+
+        Args:
+            new_settings (Dict[str, Any]): The new settings to update.
+        """
+        cl.user_session.set("llm_settings", new_settings)
+        await self.inform_llm_settings()
+        await self.setup_llm()
+
+    async def make_llm_settings_widgets(self, config=None):
+        """
+        Create and send the widgets for LLM settings configuration.
+
+        Args:
+            config: The configuration to use for setting up the widgets.
+        """
+        config = config or self.config
+        await cl.ChatSettings(
+            [
+                cl.input_widget.Select(
+                    id="chat_model",
+                    label="Model Name (Default GPT-3)",
+                    values=["local_llm", "gpt-3.5-turbo-1106", "gpt-4", "gpt-4o-mini"],
+                    initial_index=[
+                        "local_llm",
+                        "gpt-3.5-turbo-1106",
+                        "gpt-4",
+                        "gpt-4o-mini",
+                    ].index(config["llm_params"]["llm_loader"]),
+                ),
+                cl.input_widget.Select(
+                    id="retriever_method",
+                    label="Retriever (Default FAISS)",
+                    values=["FAISS", "Chroma", "RAGatouille", "RAPTOR"],
+                    initial_index=["FAISS", "Chroma", "RAGatouille", "RAPTOR"].index(
+                        config["vectorstore"]["db_option"]
+                    ),
+                ),
+                cl.input_widget.Slider(
+                    id="memory_window",
+                    label="Memory Window (Default 3)",
+                    initial=3,
+                    min=0,
+                    max=10,
+                    step=1,
+                ),
+                cl.input_widget.Switch(
+                    id="view_sources", label="View Sources", initial=False
+                ),
+                cl.input_widget.Switch(
+                    id="stream_response",
+                    label="Stream response",
+                    initial=config["llm_params"]["stream"],
+                ),
+                cl.input_widget.Switch(
+                    id="follow_up_questions",
+                    label="Generate follow up questions",
+                    initial=False,
+                ),
+                cl.input_widget.Select(
+                    id="llm_style",
+                    label="Type of Conversation (Default Normal)",
+                    values=["Normal", "ELI5"],
+                    initial_index=0,
+                ),
+            ]
+        ).send()
+
+    @no_type_check
+    async def inform_llm_settings(self):
+        """
+        Inform the user about the updated LLM settings and display them as a message.
+        """
+        llm_settings: Dict[str, Any] = cl.user_session.get("llm_settings", {})
+        llm_tutor = cl.user_session.get("llm_tutor")
+        settings_dict = {
+            "model": llm_settings.get("chat_model"),
+            "retriever": llm_settings.get("retriever_method"),
+            "memory_window": llm_settings.get("memory_window"),
+            "num_docs_in_db": (
+                len(llm_tutor.vector_db)
+                if llm_tutor and hasattr(llm_tutor, "vector_db")
+                else 0
+            ),
+            "view_sources": llm_settings.get("view_sources"),
+            "follow_up_questions": llm_settings.get("follow_up_questions"),
+        }
+        await cl.Message(
+            author=SYSTEM,
+            content="LLM settings have been updated. You can continue with your Query!",
+            elements=[
+                cl.Text(
+                    name="settings",
+                    display="side",
+                    content=json.dumps(settings_dict, indent=4),
+                    language="json",
+                ),
+            ],
+        ).send()
+
+    async def set_starters(self):
+        """
+        Set starter messages for the chatbot.
+        """
+        # Return Starters only if the chat is new
 
+        try:
+            thread = cl_data._data_layer.get_thread(
+                cl.context.session.thread_id
+            )  # see if the thread has any steps
+            if thread.steps or len(thread.steps) > 0:
+                return None
+        except:
+            return [
+                cl.Starter(
+                    label="recording on CNNs?",
+                    message="Where can I find the recording for the lecture on Transformers?",
+                    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 didn't understand the math behind backprop, could you explain it?",
+                    icon="/public/acastusphoton-svgrepo-com.svg",
+                ),
+            ]
+
+    def rename(self, orig_author: str):
+        """
+        Rename the original author to a more user-friendly name.
+
+        Args:
+            orig_author (str): The original author's name.
+
+        Returns:
+            str: The renamed author.
+        """
+        rename_dict = {"Chatbot": "AI Tutor"}
+        return rename_dict.get(orig_author, orig_author)
+
+    async def start(self, config=None):
+        """
+        Start the chatbot, initialize settings widgets,
+        and display and load previous conversation if chat logging is enabled.
+        """
+
+        start_time = time.time()
+
+        self.config = (
+            await self._load_config() if config is None else config
+        )  # Reload the configuration on chat resume
+
+        await self.make_llm_settings_widgets(self.config)  # Reload the settings widgets
+
+        await self.make_llm_settings_widgets(self.config)
+        user = cl.user_session.get("user")
+        self.user = {
+            "user_id": user.identifier,
+            "session_id": cl.context.session.thread_id,
+        }
+
+        memory = cl.user_session.get("memory", [])
+
+        cl.user_session.set("user", self.user)
+        self.llm_tutor = LLMTutor(self.config, user=self.user)
+
+        self.chain = self.llm_tutor.qa_bot(
+            memory=memory,
+            callbacks=(
+                [cl.LangchainCallbackHandler()]
+                if cl_data._data_layer and self.config["chat_logging"]["callbacks"]
+                else None
+            ),
+        )
+        self.question_generator = self.llm_tutor.question_generator
+        cl.user_session.set("llm_tutor", self.llm_tutor)
+        cl.user_session.set("chain", self.chain)
+
+        print("Time taken to start LLM: ", time.time() - start_time)
+
+    async def stream_response(self, response):
+        """
+        Stream the response from the LLM.
+
+        Args:
+            response: The response from the LLM.
+        """
+        msg = cl.Message(content="")
+        await msg.send()
+
+        output = {}
+        for chunk in response:
+            if "answer" in chunk:
+                await msg.stream_token(chunk["answer"])
+
+            for key in chunk:
+                if key not in output:
+                    output[key] = chunk[key]
+                else:
+                    output[key] += chunk[key]
+        return output
+
+    async def main(self, message):
+        """
+        Process and Display the Conversation.
+
+        Args:
+            message: The incoming chat message.
+        """
+
+        start_time = time.time()
+
+        chain = cl.user_session.get("chain")
+
+        llm_settings = cl.user_session.get("llm_settings", {})
+        view_sources = llm_settings.get("view_sources", False)
+        stream = llm_settings.get("stream_response", False)
+        steam = False  # Fix streaming
+        user_query_dict = {"input": message.content}
+        # Define the base configuration
+        chain_config = {
+            "configurable": {
+                "user_id": self.user["user_id"],
+                "conversation_id": self.user["session_id"],
+                "memory_window": self.config["llm_params"]["memory_window"],
+            }
+        }
+
+        if stream:
+            res = chain.stream(user_query=user_query_dict, config=chain_config)
+            res = await self.stream_response(res)
         else:
-            pass
+            res = await chain.invoke(
+                user_query=user_query_dict,
+                config=chain_config,
+            )
+
+        answer = res.get("answer", res.get("result"))
+
+        answer_with_sources, source_elements, sources_dict = get_sources(
+            res, answer, stream=stream, view_sources=view_sources
+        )
+        answer_with_sources = answer_with_sources.replace("$$", "$")
+
+        print("Time taken to process the message: ", time.time() - start_time)
+
+        actions = []
 
-    llm_tutor = LLMTutor(config, logger=logger)
+        if self.config["llm_params"]["generate_follow_up"]:
+            start_time = time.time()
+            list_of_questions = self.question_generator.generate_questions(
+                query=user_query_dict["input"],
+                response=answer,
+                chat_history=res.get("chat_history"),
+                context=res.get("context"),
+            )
 
-    chain = llm_tutor.qa_bot()
-    # msg = cl.Message(content=f"Starting the bot {chat_profile}...")
-    # await msg.send()
-    # msg.content = opening_message
-    # await msg.update()
+            for question in list_of_questions:
 
-    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)
+                actions.append(
+                    cl.Action(
+                        name="follow up question",
+                        value="example_value",
+                        description=question,
+                        label=question,
+                    )
+                )
 
+            print("Time taken to generate questions: ", time.time() - start_time)
 
-@cl.on_chat_end
-async def on_chat_end():
-    await cl.Message(content="Sorry, I have to go now. Goodbye!").send()
+        await cl.Message(
+            content=answer_with_sources,
+            elements=source_elements,
+            author=LLM,
+            actions=actions,
+            metadata=self.config,
+        ).send()
 
+    async def on_chat_resume(self, thread: ThreadDict):
+        thread_config = None
+        steps = thread["steps"]
+        k = self.config["llm_params"][
+            "memory_window"
+        ]  # on resume, alwyas use the default memory window
+        conversation_list = get_history_chat_resume(steps, k, SYSTEM, LLM)
+        thread_config = get_last_config(
+            steps
+        )  # TODO: Returns None for now - which causes config to be reloaded with default values
+        cl.user_session.set("memory", conversation_list)
+        await self.start(config=thread_config)
 
-@cl.on_message
-async def main(message):
-    global logger
-    user = cl.user_session.get("user")
-    chain = cl.user_session.get("chain")
+    @cl.oauth_callback
+    def auth_callback(
+        provider_id: str,
+        token: str,
+        raw_user_data: Dict[str, str],
+        default_user: cl.User,
+    ) -> Optional[cl.User]:
+        return default_user
 
-    counter = cl.user_session.get("counter")
-    counter += 1
-    cl.user_session.set("counter", counter)
+    async def on_follow_up(self, action: cl.Action):
+        message = await cl.Message(
+            content=action.description,
+            type="user_message",
+            author=self.user["user_id"],
+        ).send()
+        await self.main(message)
 
-    # 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
+chatbot = Chatbot(config=config)
 
-    processor = cl.user_session.get("chat_processor")
-    res = await processor.rag(message.content, chain, cb)
-    try:
-        answer = res["answer"]
-    except:
-        answer = res["result"]
 
-    answer_with_sources, source_elements, sources_dict = get_sources(res, answer)
-    processor._process(message.content, answer, sources_dict)
+async def start_app():
+    cl_data._data_layer = await setup_data_layer()
+    chatbot.literal_client = cl_data._data_layer.client if cl_data._data_layer else None
+    cl.set_starters(chatbot.set_starters)
+    cl.author_rename(chatbot.rename)
+    cl.on_chat_start(chatbot.start)
+    cl.on_chat_resume(chatbot.on_chat_resume)
+    cl.on_message(chatbot.main)
+    cl.on_settings_update(chatbot.update_llm)
+    cl.action_callback("follow up question")(chatbot.on_follow_up)
 
-    answer_with_sources = answer_with_sources.replace("$$", "$")
 
-    await cl.Message(content=answer_with_sources, elements=source_elements).send()
+asyncio.run(start_app())

code/modules/chat/base.py

@@ -0,0 +1,13 @@
+class BaseRAG:
+    """
+    Base class for RAG chatbot.
+    """
+
+    def __init__():
+        pass
+
+    def invoke():
+        """
+        Invoke the RAG chatbot.
+        """
+        pass

code/modules/chat/chat_model_loader.py

@@ -1,4 +1,4 @@
-from langchain_community.chat_models import ChatOpenAI
+from langchain_openai import ChatOpenAI
 from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
 from transformers import AutoTokenizer, TextStreamer
 from langchain_community.llms import LlamaCpp
@@ -9,6 +9,7 @@ from pathlib import Path
 from huggingface_hub import hf_hub_download
 from langchain.callbacks.manager import CallbackManager
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from modules.config.constants import LLAMA_PATH
 
 
 class ChatModelLoader:
@@ -20,20 +21,24 @@ class ChatModelLoader:
         hf_hub_download(
             repo_id=self.config["llm_params"]["local_llm_params"]["repo_id"],
             filename=self.config["llm_params"]["local_llm_params"]["filename"],
-            cache_dir=model_cache_path
-            )
+            cache_dir=model_cache_path,
+        )
         return str(list(Path(model_cache_path).glob("*/snapshots/*/*.gguf"))[0])
 
     def load_chat_model(self):
-        if self.config["llm_params"]["llm_loader"] == "openai":
-            llm = ChatOpenAI(
-                model_name=self.config["llm_params"]["openai_params"]["model"]
-            )
+        if self.config["llm_params"]["llm_loader"] in [
+            "gpt-3.5-turbo-1106",
+            "gpt-4",
+            "gpt-4o-mini",
+        ]:
+            llm = ChatOpenAI(model_name=self.config["llm_params"]["llm_loader"])
         elif self.config["llm_params"]["llm_loader"] == "local_llm":
             n_batch = 512  # Should be between 1 and n_ctx, consider the amount of VRAM in your GPU.
-            model_path = self._verify_model_cache(self.config["llm_params"]["local_llm_params"]["model"])
+            model_path = self._verify_model_cache(
+                self.config["llm_params"]["local_llm_params"]["model"]
+            )
             llm = LlamaCpp(
-                model_path=model_path,
+                model_path=LLAMA_PATH,
                 n_batch=n_batch,
                 n_ctx=2048,
                 f16_kv=True,
@@ -44,5 +49,7 @@ class ChatModelLoader:
                 ],
             )
         else:
-            raise ValueError("Invalid LLM Loader")
+            raise ValueError(
+                f"Invalid LLM Loader: {self.config['llm_params']['llm_loader']}"
+            )
         return llm

code/modules/chat/helpers.py

@@ -1,13 +1,12 @@
-from modules.config.constants import *
+from modules.config.prompts import prompts
 import chainlit as cl
-from langchain_core.prompts import PromptTemplate
 
 
-def get_sources(res, answer):
+def get_sources(res, answer, stream=True, view_sources=False):
     source_elements = []
     source_dict = {}  # Dictionary to store URL elements
 
-    for idx, source in enumerate(res["source_documents"]):
+    for idx, source in enumerate(res["context"]):
         source_metadata = source.metadata
         url = source_metadata.get("source", "N/A")
         score = source_metadata.get("score", "N/A")
@@ -36,69 +35,135 @@ def get_sources(res, answer):
         else:
             source_dict[url_name]["text"] += f"\n\n{source.page_content}"
 
-    # First, display the answer
-    full_answer = "**Answer:**\n"
-    full_answer += answer
+    full_answer = ""  # Not to include the answer again if streaming
 
-    # 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"
+    if not stream:  # First, display the answer if not streaming
+        full_answer = "**Answer:**\n"
+        full_answer += answer
 
-        name = f"Source {idx + 1} Text\n"
-        full_answer += name
-        source_elements.append(
-            cl.Text(name=name, content=source_data["text"], display="side")
-        )
+    if view_sources:
 
-        # 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",
-            )
-        )
+        # Then, display the sources
+        # check if the answer has sources
+        if len(source_dict) == 0:
+            full_answer += "\n\n**No sources found.**"
+            return full_answer, source_elements, source_dict
+        else:
+            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"],
+def get_prompt(config, prompt_type):
+    llm_params = config["llm_params"]
+    llm_loader = llm_params["llm_loader"]
+    use_history = llm_params["use_history"]
+    llm_style = llm_params["llm_style"].lower()
+
+    if prompt_type == "qa":
+        if llm_loader == "local_llm":
+            if use_history:
+                return prompts["tiny_llama"]["prompt_with_history"]
+            else:
+                return prompts["tiny_llama"]["prompt_no_history"]
+        else:
+            if use_history:
+                return prompts["openai"]["prompt_with_history"][llm_style]
+            else:
+                return prompts["openai"]["prompt_no_history"]
+    elif prompt_type == "rephrase":
+        return prompts["openai"]["rephrase_prompt"]
+
+
+def get_history_chat_resume(steps, k, SYSTEM, LLM):
+    conversation_list = []
+    count = 0
+    for step in reversed(steps):
+        if step["name"] not in [SYSTEM]:
+            if step["type"] == "user_message":
+                conversation_list.append(
+                    {"type": "user_message", "content": step["output"]}
+                )
+            elif step["type"] == "assistant_message":
+                if step["name"] == LLM:
+                    conversation_list.append(
+                        {"type": "ai_message", "content": step["output"]}
+                    )
+            else:
+                raise ValueError("Invalid message type")
+        count += 1
+        if count >= 2 * k:  # 2 * k to account for both user and assistant messages
+            break
+    conversation_list = conversation_list[::-1]
+    return conversation_list
+
+
+def get_history_setup_llm(memory_list):
+    conversation_list = []
+    for message in memory_list:
+        message_dict = message.to_dict() if hasattr(message, "to_dict") else message
+
+        # Check if the type attribute is present as a key or attribute
+        message_type = (
+            message_dict.get("type", None)
+            if isinstance(message_dict, dict)
+            else getattr(message, "type", None)
         )
-    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"],
+
+        # Check if content is present as a key or attribute
+        message_content = (
+            message_dict.get("content", None)
+            if isinstance(message_dict, dict)
+            else getattr(message, "content", None)
         )
-    return prompt
+
+        if message_type in ["ai", "ai_message"]:
+            conversation_list.append({"type": "ai_message", "content": message_content})
+        elif message_type in ["human", "user_message"]:
+            conversation_list.append(
+                {"type": "user_message", "content": message_content}
+            )
+        else:
+            raise ValueError("Invalid message type")
+
+    return conversation_list
+
+
+def get_last_config(steps):
+    # TODO: Implement this function
+    return None

code/modules/chat/langchain/langchain_rag.py

@@ -0,0 +1,274 @@
+from langchain_core.prompts import ChatPromptTemplate
+
+from modules.chat.langchain.utils import *
+from langchain.memory import ChatMessageHistory
+from modules.chat.base import BaseRAG
+from langchain_core.prompts import PromptTemplate
+from langchain.memory import (
+    ConversationBufferWindowMemory,
+    ConversationSummaryBufferMemory,
+)
+
+import chainlit as cl
+from langchain_community.chat_models import ChatOpenAI
+
+
+class Langchain_RAG_V1(BaseRAG):
+
+    def __init__(
+        self,
+        llm,
+        memory,
+        retriever,
+        qa_prompt: str,
+        rephrase_prompt: str,
+        config: dict,
+        callbacks=None,
+    ):
+        """
+        Initialize the Langchain_RAG class.
+
+        Args:
+            llm (LanguageModelLike): The language model instance.
+            memory (BaseChatMessageHistory): The chat message history instance.
+            retriever (BaseRetriever): The retriever instance.
+            qa_prompt (str): The QA prompt string.
+            rephrase_prompt (str): The rephrase prompt string.
+        """
+        self.llm = llm
+        self.config = config
+        # self.memory = self.add_history_from_list(memory)
+        self.memory = ConversationBufferWindowMemory(
+            k=self.config["llm_params"]["memory_window"],
+            memory_key="chat_history",
+            return_messages=True,
+            output_key="answer",
+            max_token_limit=128,
+        )
+        self.retriever = retriever
+        self.qa_prompt = qa_prompt
+        self.rephrase_prompt = rephrase_prompt
+        self.store = {}
+
+        self.qa_prompt = PromptTemplate(
+            template=self.qa_prompt,
+            input_variables=["context", "chat_history", "input"],
+        )
+
+        self.rag_chain = CustomConversationalRetrievalChain.from_llm(
+            llm=llm,
+            chain_type="stuff",
+            retriever=retriever,
+            return_source_documents=True,
+            memory=self.memory,
+            combine_docs_chain_kwargs={"prompt": self.qa_prompt},
+            response_if_no_docs_found="No context found",
+        )
+
+    def add_history_from_list(self, history_list):
+        """
+        TODO: Add messages from a list to the chat history.
+        """
+        history = []
+
+        return history
+
+    async def invoke(self, user_query, config):
+        """
+        Invoke the chain.
+
+        Args:
+            kwargs: The input variables.
+
+        Returns:
+            dict: The output variables.
+        """
+        res = await self.rag_chain.acall(user_query["input"])
+        return res
+
+
+class QuestionGenerator:
+    """
+    Generate a question from the LLMs response and users input and past conversations.
+    """
+
+    def __init__(self):
+        pass
+
+    def generate_questions(self, query, response, chat_history, context):
+        questions = return_questions(query, response, chat_history, context)
+        return questions
+
+
+class Langchain_RAG_V2(BaseRAG):
+    def __init__(
+        self,
+        llm,
+        memory,
+        retriever,
+        qa_prompt: str,
+        rephrase_prompt: str,
+        config: dict,
+        callbacks=None,
+    ):
+        """
+        Initialize the Langchain_RAG class.
+
+        Args:
+            llm (LanguageModelLike): The language model instance.
+            memory (BaseChatMessageHistory): The chat message history instance.
+            retriever (BaseRetriever): The retriever instance.
+            qa_prompt (str): The QA prompt string.
+            rephrase_prompt (str): The rephrase prompt string.
+        """
+        self.llm = llm
+        self.memory = self.add_history_from_list(memory)
+        self.retriever = retriever
+        self.qa_prompt = qa_prompt
+        self.rephrase_prompt = rephrase_prompt
+        self.store = {}
+
+        # Contextualize question prompt
+        contextualize_q_system_prompt = rephrase_prompt or (
+            "Given a chat history and the latest user question "
+            "which might reference context in the chat history, "
+            "formulate a standalone question which can be understood "
+            "without the chat history. Do NOT answer the question, just "
+            "reformulate it if needed and otherwise return it as is."
+        )
+        self.contextualize_q_prompt = ChatPromptTemplate.from_template(
+            contextualize_q_system_prompt
+        )
+
+        # History-aware retriever
+        self.history_aware_retriever = create_history_aware_retriever(
+            self.llm, self.retriever, self.contextualize_q_prompt
+        )
+
+        # Answer question prompt
+        qa_system_prompt = qa_prompt or (
+            "You are an assistant for question-answering tasks. Use "
+            "the following pieces of retrieved context to answer the "
+            "question. If you don't know the answer, just say that you "
+            "don't know. Use three sentences maximum and keep the answer "
+            "concise."
+            "\n\n"
+            "{context}"
+        )
+        self.qa_prompt_template = ChatPromptTemplate.from_template(qa_system_prompt)
+
+        # Question-answer chain
+        self.question_answer_chain = create_stuff_documents_chain(
+            self.llm, self.qa_prompt_template
+        )
+
+        # Final retrieval chain
+        self.rag_chain = create_retrieval_chain(
+            self.history_aware_retriever, self.question_answer_chain
+        )
+
+        self.rag_chain = CustomRunnableWithHistory(
+            self.rag_chain,
+            get_session_history=self.get_session_history,
+            input_messages_key="input",
+            history_messages_key="chat_history",
+            output_messages_key="answer",
+            history_factory_config=[
+                ConfigurableFieldSpec(
+                    id="user_id",
+                    annotation=str,
+                    name="User ID",
+                    description="Unique identifier for the user.",
+                    default="",
+                    is_shared=True,
+                ),
+                ConfigurableFieldSpec(
+                    id="conversation_id",
+                    annotation=str,
+                    name="Conversation ID",
+                    description="Unique identifier for the conversation.",
+                    default="",
+                    is_shared=True,
+                ),
+                ConfigurableFieldSpec(
+                    id="memory_window",
+                    annotation=int,
+                    name="Number of Conversations",
+                    description="Number of conversations to consider for context.",
+                    default=1,
+                    is_shared=True,
+                ),
+            ],
+        )
+
+        if callbacks is not None:
+            self.rag_chain = self.rag_chain.with_config(callbacks=callbacks)
+
+    def get_session_history(
+        self, user_id: str, conversation_id: str, memory_window: int
+    ) -> BaseChatMessageHistory:
+        """
+        Get the session history for a user and conversation.
+
+        Args:
+            user_id (str): The user identifier.
+            conversation_id (str): The conversation identifier.
+            memory_window (int): The number of conversations to consider for context.
+
+        Returns:
+            BaseChatMessageHistory: The chat message history.
+        """
+        if (user_id, conversation_id) not in self.store:
+            self.store[(user_id, conversation_id)] = InMemoryHistory()
+            self.store[(user_id, conversation_id)].add_messages(
+                self.memory.messages
+            )  # add previous messages to the store. Note: the store is in-memory.
+        return self.store[(user_id, conversation_id)]
+
+    async def invoke(self, user_query, config, **kwargs):
+        """
+        Invoke the chain.
+
+        Args:
+            kwargs: The input variables.
+
+        Returns:
+            dict: The output variables.
+        """
+        res = await self.rag_chain.ainvoke(user_query, config, **kwargs)
+        res["rephrase_prompt"] = self.rephrase_prompt
+        res["qa_prompt"] = self.qa_prompt
+        return res
+
+    def stream(self, user_query, config):
+        res = self.rag_chain.stream(user_query, config)
+        return res
+
+    def add_history_from_list(self, conversation_list):
+        """
+        Add messages from a list to the chat history.
+
+        Args:
+            messages (list): The list of messages to add.
+        """
+        history = ChatMessageHistory()
+
+        for idx, message in enumerate(conversation_list):
+            message_type = (
+                message.get("type", None)
+                if isinstance(message, dict)
+                else getattr(message, "type", None)
+            )
+
+            message_content = (
+                message.get("content", None)
+                if isinstance(message, dict)
+                else getattr(message, "content", None)
+            )
+
+            if message_type in ["human", "user_message"]:
+                history.add_user_message(message_content)
+            elif message_type in ["ai", "ai_message"]:
+                history.add_ai_message(message_content)
+
+        return history

code/modules/chat/langchain/utils.py

@@ -0,0 +1,340 @@
+from typing import Any, Dict, List, Union, Tuple, Optional
+from langchain_core.messages import (
+    BaseMessage,
+    AIMessage,
+    FunctionMessage,
+    HumanMessage,
+)
+
+from langchain_core.prompts.base import BasePromptTemplate, format_document
+from langchain_core.prompts.chat import MessagesPlaceholder
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.output_parsers.base import BaseOutputParser
+from langchain_core.retrievers import BaseRetriever, RetrieverOutput
+from langchain_core.language_models import LanguageModelLike
+from langchain_core.runnables import Runnable, RunnableBranch, RunnablePassthrough
+from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain_core.runnables.utils import ConfigurableFieldSpec
+from langchain_core.chat_history import BaseChatMessageHistory
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain.chains.combine_documents.base import (
+    DEFAULT_DOCUMENT_PROMPT,
+    DEFAULT_DOCUMENT_SEPARATOR,
+    DOCUMENTS_KEY,
+    BaseCombineDocumentsChain,
+    _validate_prompt,
+)
+from langchain.chains.llm import LLMChain
+from langchain_core.callbacks import Callbacks
+from langchain_core.documents import Document
+
+
+CHAT_TURN_TYPE = Union[Tuple[str, str], BaseMessage]
+
+from langchain_core.runnables.config import RunnableConfig
+from langchain_core.messages import BaseMessage
+
+
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_community.chat_models import ChatOpenAI
+
+from langchain.chains import RetrievalQA, ConversationalRetrievalChain
+from langchain_core.callbacks.manager import AsyncCallbackManagerForChainRun
+
+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
+
+
+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: '{input}'"
+            )
+
+            prompt = ChatPromptTemplate.from_messages(
+                [
+                    ("system", system),
+                    ("human", "{input}, {chat_history_str}"),
+                ]
+            )
+            llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
+            step_back = prompt | llm | StrOutputParser()
+            new_question = step_back.invoke(
+                {"input": 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: {input}\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
+        output["context"] = output["source_documents"]
+        return output
+
+
+class CustomRunnableWithHistory(RunnableWithMessageHistory):
+
+    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 _aenter_history(
+        self, input: Any, config: RunnableConfig
+    ) -> List[BaseMessage]:
+        """
+        Get the last k conversations from the message history.
+
+        Args:
+            input (Any): The input data.
+            config (RunnableConfig): The runnable configuration.
+
+        Returns:
+            List[BaseMessage]: The last k conversations.
+        """
+        hist: BaseChatMessageHistory = config["configurable"]["message_history"]
+        messages = (await hist.aget_messages()).copy()
+        if not self.history_messages_key:
+            # return all messages
+            input_val = (
+                input if not self.input_messages_key else input[self.input_messages_key]
+            )
+            messages += self._get_input_messages(input_val)
+
+        # return last k conversations
+        if config["configurable"]["memory_window"] == 0:  # if k is 0, return empty list
+            messages = []
+        else:
+            messages = messages[-2 * config["configurable"]["memory_window"] :]
+
+        messages = self._get_chat_history(messages)
+
+        return messages
+
+
+class InMemoryHistory(BaseChatMessageHistory, BaseModel):
+    """In-memory implementation of chat message history."""
+
+    messages: List[BaseMessage] = Field(default_factory=list)
+
+    def add_messages(self, messages: List[BaseMessage]) -> None:
+        """Add a list of messages to the store."""
+        self.messages.extend(messages)
+
+    def clear(self) -> None:
+        """Clear the message history."""
+        self.messages = []
+
+    def __len__(self) -> int:
+        """Return the number of messages."""
+        return len(self.messages)
+
+
+def create_history_aware_retriever(
+    llm: LanguageModelLike,
+    retriever: BaseRetriever,
+    prompt: BasePromptTemplate,
+) -> Runnable[Dict[str, Any], RetrieverOutput]:
+    """Create a chain that takes conversation history and returns documents."""
+    if "input" not in prompt.input_variables:
+        raise ValueError(
+            "Expected `input` to be a prompt variable, "
+            f"but got {prompt.input_variables}"
+        )
+
+    retrieve_documents = RunnableBranch(
+        (
+            lambda x: not x["chat_history"],
+            (lambda x: x["input"]) | retriever,
+        ),
+        prompt | llm | StrOutputParser() | retriever,
+    ).with_config(run_name="chat_retriever_chain")
+
+    return retrieve_documents
+
+
+def create_stuff_documents_chain(
+    llm: LanguageModelLike,
+    prompt: BasePromptTemplate,
+    output_parser: Optional[BaseOutputParser] = None,
+    document_prompt: Optional[BasePromptTemplate] = None,
+    document_separator: str = DEFAULT_DOCUMENT_SEPARATOR,
+) -> Runnable[Dict[str, Any], Any]:
+    """Create a chain for passing a list of Documents to a model."""
+    _validate_prompt(prompt)
+    _document_prompt = document_prompt or DEFAULT_DOCUMENT_PROMPT
+    _output_parser = output_parser or StrOutputParser()
+
+    def format_docs(inputs: dict) -> str:
+        return document_separator.join(
+            format_document(doc, _document_prompt) for doc in inputs[DOCUMENTS_KEY]
+        )
+
+    return (
+        RunnablePassthrough.assign(**{DOCUMENTS_KEY: format_docs}).with_config(
+            run_name="format_inputs"
+        )
+        | prompt
+        | llm
+        | _output_parser
+    ).with_config(run_name="stuff_documents_chain")
+
+
+def create_retrieval_chain(
+    retriever: Union[BaseRetriever, Runnable[dict, RetrieverOutput]],
+    combine_docs_chain: Runnable[Dict[str, Any], str],
+) -> Runnable:
+    """Create retrieval chain that retrieves documents and then passes them on."""
+    if not isinstance(retriever, BaseRetriever):
+        retrieval_docs = retriever
+    else:
+        retrieval_docs = (lambda x: x["input"]) | retriever
+
+    retrieval_chain = (
+        RunnablePassthrough.assign(
+            context=retrieval_docs.with_config(run_name="retrieve_documents"),
+        ).assign(answer=combine_docs_chain)
+    ).with_config(run_name="retrieval_chain")
+
+    return retrieval_chain
+
+
+def return_questions(query, response, chat_history_str, context):
+
+    system = (
+        "You are someone that suggests a question based on the student's input and chat history. "
+        "Generate a question that is relevant to the student's input and chat history. "
+        "Incorporate relevant details from the chat history to make the question clearer and more specific. "
+        "Chat history: \n{chat_history_str}\n"
+        "Use the context to generate a question that is relevant to the student's input and chat history: Context: {context}"
+        "Generate 3 short and concise questions from the students voice based on the following input and response: "
+        "The 3 short and concise questions should be sperated by dots. Example: 'What is the capital of France?...What is the population of France?...What is the currency of France?'"
+        "User Query: {query}"
+        "AI Response: {response}"
+        "The 3 short and concise questions seperated by dots (...) are:"
+    )
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system),
+            ("human", "{chat_history_str}, {context}, {query}, {response}"),
+        ]
+    )
+    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+    question_generator = prompt | llm | StrOutputParser()
+    new_questions = question_generator.invoke(
+        {
+            "chat_history_str": chat_history_str,
+            "context": context,
+            "query": query,
+            "response": response,
+        }
+    )
+
+    list_of_questions = new_questions.split("...")
+    return list_of_questions

code/modules/chat/llm_tutor.py

@@ -1,211 +1,163 @@
-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
+from modules.chat.langchain.langchain_rag import (
+    Langchain_RAG_V1,
+    Langchain_RAG_V2,
+    QuestionGenerator,
+)
 
 
 class LLMTutor:
-    def __init__(self, config, logger=None):
+    def __init__(self, config, user, logger=None):
+        """
+        Initialize the LLMTutor class.
+
+        Args:
+            config (dict): Configuration dictionary.
+            user (str): User identifier.
+            logger (Logger, optional): Logger instance. Defaults to None.
+        """
         self.config = config
         self.llm = self.load_llm()
+        self.user = user
         self.logger = logger
-        self.vector_db = VectorStoreManager(config, logger=self.logger)
+        self.vector_db = VectorStoreManager(config, logger=self.logger).load_database()
+        self.qa_prompt = get_prompt(config, "qa")  # Initialize qa_prompt
+        self.rephrase_prompt = get_prompt(
+            config, "rephrase"
+        )  # Initialize rephrase_prompt
         if self.config["vectorstore"]["embedd_files"]:
             self.vector_db.create_database()
             self.vector_db.save_database()
 
-    def set_custom_prompt(self):
+    def update_llm(self, old_config, new_config):
         """
-        Prompt template for QA retrieval for each vectorstore
+        Update the LLM and VectorStoreManager based on new configuration.
+
+        Args:
+            new_config (dict): New configuration dictionary.
         """
-        prompt = get_prompt(self.config)
-        # prompt = QA_PROMPT
+        changes = self.get_config_changes(old_config, new_config)
+
+        if "llm_params.llm_loader" in changes:
+            self.llm = self.load_llm()  # Reinitialize LLM if chat_model changes
 
-        return prompt
+        if "vectorstore.db_option" in changes:
+            self.vector_db = VectorStoreManager(
+                self.config, logger=self.logger
+            ).load_database()  # Reinitialize VectorStoreManager if vectorstore changes
+            if self.config["vectorstore"]["embedd_files"]:
+                self.vector_db.create_database()
+                self.vector_db.save_database()
 
-    # Retrieval QA Chain
-    def retrieval_qa_chain(self, llm, prompt, db):
+        if "llm_params.llm_style" in changes:
+            self.qa_prompt = get_prompt(
+                self.config, "qa"
+            )  # Update qa_prompt if ELI5 changes
 
+    def get_config_changes(self, old_config, new_config):
+        """
+        Get the changes between the old and new configuration.
+
+        Args:
+            old_config (dict): Old configuration dictionary.
+            new_config (dict): New configuration dictionary.
+
+        Returns:
+            dict: Dictionary containing the changes.
+        """
+        changes = {}
+
+        def compare_dicts(old, new, parent_key=""):
+            for key in new:
+                full_key = f"{parent_key}.{key}" if parent_key else key
+                if isinstance(new[key], dict) and isinstance(old.get(key), dict):
+                    compare_dicts(old.get(key, {}), new[key], full_key)
+                elif old.get(key) != new[key]:
+                    changes[full_key] = (old.get(key), new[key])
+            # Include keys that are in old but not in new
+            for key in old:
+                if key not in new:
+                    full_key = f"{parent_key}.{key}" if parent_key else key
+                    changes[full_key] = (old[key], None)
+
+        compare_dicts(old_config, new_config)
+        return changes
+
+    def retrieval_qa_chain(
+        self, llm, qa_prompt, rephrase_prompt, db, memory=None, callbacks=None
+    ):
+        """
+        Create a Retrieval QA Chain.
+
+        Args:
+            llm (LLM): The language model instance.
+            qa_prompt (str): The QA prompt string.
+            rephrase_prompt (str): The rephrase prompt string.
+            db (VectorStore): The vector store instance.
+            memory (Memory, optional): Memory instance. Defaults to None.
+
+        Returns:
+            Chain: The retrieval QA chain instance.
+        """
         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(
+        if self.config["llm_params"]["llm_arch"] == "langchain":
+            self.qa_chain = Langchain_RAG_V2(
                 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",
+                retriever=retriever,
+                qa_prompt=qa_prompt,
+                rephrase_prompt=rephrase_prompt,
+                config=self.config,
+                callbacks=callbacks,
             )
+
+            self.question_generator = QuestionGenerator()
         else:
-            qa_chain = RetrievalQA.from_chain_type(
-                llm=llm,
-                chain_type="stuff",
-                retriever=retriever,
-                return_source_documents=True,
-                chain_type_kwargs={"prompt": prompt},
+            raise ValueError(
+                f"Invalid LLM Architecture: {self.config['llm_params']['llm_arch']}"
             )
-        return qa_chain
+        return self.qa_chain
 
-    # Loading the model
     def load_llm(self):
+        """
+        Load the language model.
+
+        Returns:
+            LLM: The loaded language model instance.
+        """
         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()
+    def qa_bot(self, memory=None, callbacks=None):
+        """
+        Create a QA bot instance.
+
+        Args:
+            memory (Memory, optional): Memory instance. Defaults to None.
+            qa_prompt (str, optional): QA prompt string. Defaults to None.
+            rephrase_prompt (str, optional): Rephrase prompt string. Defaults to None.
+
+        Returns:
+            Chain: The QA bot chain instance.
+        """
+        # sanity check to see if there are any documents in the database
+        if len(self.vector_db) == 0:
+            raise ValueError(
+                "No documents in the database. Populate the database first."
+            )
+
         qa = self.retrieval_qa_chain(
-            self.llm, qa_prompt, db
-        )  # TODO:  PROMPT is overwritten in CustomConversationalRetrievalChain
+            self.llm,
+            self.qa_prompt,
+            self.rephrase_prompt,
+            self.vector_db,
+            memory,
+            callbacks=callbacks,
+        )
 
         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

@@ -1,12 +0,0 @@
-# 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

@@ -1,30 +0,0 @@
-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

@@ -1,37 +1,44 @@
-from literalai import LiteralClient
-import os
-from .base import ChatProcessorBase
+from chainlit.data import ChainlitDataLayer, queue_until_user_message
 
 
-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}")
+# update custom methods here (Ref: https://github.com/Chainlit/chainlit/blob/4b533cd53173bcc24abe4341a7108f0070d60099/backend/chainlit/data/__init__.py)
+class CustomLiteralDataLayer(ChainlitDataLayer):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
 
-    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",
-            )
+    @queue_until_user_message()
+    async def create_step(self, step_dict: "StepDict"):
+        metadata = dict(
+            step_dict.get("metadata", {}),
+            **{
+                "waitForAnswer": step_dict.get("waitForAnswer"),
+                "language": step_dict.get("language"),
+                "showInput": step_dict.get("showInput"),
+            },
+        )
 
-    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
+        step: LiteralStepDict = {
+            "createdAt": step_dict.get("createdAt"),
+            "startTime": step_dict.get("start"),
+            "endTime": step_dict.get("end"),
+            "generation": step_dict.get("generation"),
+            "id": step_dict.get("id"),
+            "parentId": step_dict.get("parentId"),
+            "name": step_dict.get("name"),
+            "threadId": step_dict.get("threadId"),
+            "type": step_dict.get("type"),
+            "tags": step_dict.get("tags"),
+            "metadata": metadata,
+        }
+        if step_dict.get("input"):
+            step["input"] = {"content": step_dict.get("input")}
+        if step_dict.get("output"):
+            step["output"] = {"content": step_dict.get("output")}
+        if step_dict.get("isError"):
+            step["error"] = step_dict.get("output")
+
+        # print("\n\n\n")
+        # print("Step: ", step)
+        # print("\n\n\n")
+
+        await self.client.api.send_steps([step])

code/modules/config/config.yml

@@ -4,7 +4,6 @@ device: 'cpu' # str [cuda, cpu]
 
 vectorstore:
   load_from_HF: True # bool
-  HF_path: "XThomasBU/Colbert_Index" # str
   embedd_files: False # bool
   data_path: '../storage/data' # str
   url_file_path: '../storage/data/urls.txt' # str
@@ -16,7 +15,7 @@ vectorstore:
   score_threshold : 0.2 # float
 
   faiss_params: # Not used as of now
-    index_path: '../vectorstores/faiss.index' # str
+    index_path: 'vectorstores/faiss.index' # str
     index_type: 'Flat' # str [Flat, HNSW, IVF]
     index_dimension: 384 # int
     index_nlist: 100 # int
@@ -26,21 +25,25 @@ vectorstore:
     index_name: "new_idx" # str
 
 llm_params: 
+  llm_arch: 'langchain' # [langchain]
   use_history: True # bool
+  generate_follow_up: False # bool
   memory_window: 3 # int
-  llm_loader: 'openai' # str [local_llm, openai]
+  llm_style: 'Normal' # str [Normal, ELI5]
+  llm_loader: 'gpt-4o-mini' # str [local_llm, gpt-3.5-turbo-1106, gpt-4, gpt-4o-mini]
   openai_params:
-    model: 'gpt-3.5-turbo-1106' # str [gpt-3.5-turbo-1106, gpt-4]
+    temperature: 0.7 # float
   local_llm_params:
-    model: 'tiny-llama'
-    temperature: 0.7
+    temperature: 0.7 # float
     repo_id: 'TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF' # HuggingFace repo id
     filename: 'tinyllama-1.1b-chat-v1.0.Q5_0.gguf' # Specific name of gguf file in the repo
-  pdf_reader: 'llama' # str [llama, pymupdf, gpt]
+  pdf_reader: 'pymupdf' # str [llama, pymupdf, gpt]
+  stream: False # bool
 
 chat_logging:
-  log_chat: False # bool
+  log_chat: True # bool
   platform: 'literalai'
+  callbacks: False # bool
 
 splitter_options:
   use_splitter: True # bool

code/modules/config/constants.py

@@ -8,83 +8,16 @@ load_dotenv()
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 LLAMA_CLOUD_API_KEY = os.getenv("LLAMA_CLOUD_API_KEY")
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
-LITERAL_API_KEY = os.getenv("LITERAL_API_KEY")
+LITERAL_API_KEY_LOGGING = os.getenv("LITERAL_API_KEY_LOGGING")
+LITERAL_API_URL = os.getenv("LITERAL_API_URL")
 
-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
-
-openai_prompt_template = """Use the following pieces of information to answer the user's question.
-You are an intelligent chatbot designed to help students with questions regarding the course. 
-Render math equations in LaTeX format between $ or $$ signs, stick to the parameter and variable icons found in your context. 
-Be sure to explain the parameters and variables in the equations.
-If you don't know the answer, just say that you don't know. 
-
-Context: {context}
-Question: {question}
-
-Only return the helpful answer below and nothing else.
-Helpful answer:
-"""
-
-openai_prompt_template_with_history = """Use the following pieces of information to answer the user's question.
-You are an intelligent chatbot designed to help students with questions regarding the course. 
-Render math equations in LaTeX format between $ or $$ signs, stick to the parameter and variable icons found in your context. 
-Be sure to explain the parameters and variables in the equations.
-If you don't know the answer, just say that you don't know, don't try to make up an answer.
-
-Use the history to answer the question if you can.
-Chat History:
-{chat_history}
-Context: {context}
-Question: {question}
-
-Only return the helpful answer below and nothing else.
-Helpful answer:
-"""
-
-tinyllama_prompt_template = """
-<|im_start|>system
-Assistant is an intelligent chatbot designed to help students with questions regarding the course. Only answer questions using the context below and if you're not sure of an answer, you can say "I don't know". Always give a brief and concise answer to the question. When asked for formulas, give a brief description of the formula and output math equations in LaTeX format between $ signs.
-
-Context:
-{context}
-<|im_end|>
-<|im_start|>user
-Question: Who is the instructor for this course?
-<|im_end|>
-<|im_start|>assistant
-The instructor for this course is Prof. Thomas Gardos.
-<|im_end|>
-<|im_start|>user
-Question: {question}
-<|im_end|>
-<|im_start|>assistant
-"""
-
-tinyllama_prompt_template_with_history = """
-<|im_start|>system
-Assistant is an intelligent chatbot designed to help students with questions regarding the course. Only answer questions using the context below and if you're not sure of an answer, you can say "I don't know". Always give a brief and concise answer to the question. Output math equations in LaTeX format between $ signs. Use the history to answer the question if you can.
-
-Chat History:
-{chat_history}
-Context:
-{context}
-<|im_end|>
-<|im_start|>user
-Question: Who is the instructor for this course?
-<|im_end|>
-<|im_start|>assistant
-The instructor for this course is Prof. Thomas Gardos.
-<|im_end|>
-<|im_start|>user
-Question: {question}
-<|im_end|>
-<|im_start|>assistant
-"""
+OAUTH_GOOGLE_CLIENT_ID = os.getenv("OAUTH_GOOGLE_CLIENT_ID")
+OAUTH_GOOGLE_CLIENT_SECRET = os.getenv("OAUTH_GOOGLE_CLIENT_SECRET")
 
+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!"
 
 # Model Paths
 
 LLAMA_PATH = "../storage/models/tinyllama"
-MISTRAL_PATH = "storage/models/mistral-7b-v0.1.Q4_K_M.gguf"
+
+RETRIEVER_HF_PATHS = {"RAGatouille": "XThomasBU/Colbert_Index"}

code/modules/config/prompts.py

@@ -0,0 +1,97 @@
+prompts = {
+    "openai": {
+        "rephrase_prompt": (
+            "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 backpropagation.'. "
+            "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}\n"
+            "Rephrase the following question only if necessary: '{input}'"
+            "Rephrased Question:'"
+        ),
+        "prompt_with_history": {
+            "normal": (
+                "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 relevant. "
+                "Render math equations in LaTeX format between $ or $$ signs, stick to the parameter and variable icons found in your context. Be sure to explain the parameters and variables in the equations."
+                "Speak in a friendly and engaging manner, like talking to a friend. Avoid sounding repetitive or robotic.\n\n"
+                "Do not get influenced by the style of conversation in the chat history. Follow the instructions given here."
+                "Chat History:\n{chat_history}\n\n"
+                "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"
+                "Student: {input}\n"
+                "AI Tutor:"
+            ),
+            "eli5": (
+                "You are an AI Tutor for the course DS598, taught by Prof. Thomas Gardos. Your job is to explain things in the simplest and most engaging way possible, just like the 'Explain Like I'm 5' (ELI5) concept."
+                "If you don't know the answer, do your best without making things up. Keep your explanations straightforward and very easy to understand."
+                "Use the chat history and context to help you, but avoid repeating past responses. Provide links from the source_file metadata when they're helpful."
+                "Use very simple language and examples to explain any math equations, and put the equations in LaTeX format between $ or $$ signs."
+                "Be friendly and engaging, like you're chatting with a young child who's curious and eager to learn. Avoid complex terms and jargon."
+                "Include simple and clear examples wherever you can to make things easier to understand."
+                "Do not get influenced by the style of conversation in the chat history. Follow the instructions given here."
+                "Chat History:\n{chat_history}\n\n"
+                "Context:\n{context}\n\n"
+                "Answer the student's question below in a friendly, simple, and engaging way, just like the ELI5 concept. Use the context and history only if they're relevant, otherwise, just have a natural conversation."
+                "Give a clear and detailed explanation with simple examples to make it easier to understand. Remember, your goal is to break down complex topics into very simple terms, just like ELI5."
+                "Student: {input}\n"
+                "AI Tutor:"
+            ),
+            "socratic": (
+                "You are an AI Tutor for the course DS598, taught by Prof. Thomas Gardos. Engage the student in a Socratic dialogue to help them discover answers on their own. Use the provided context to guide your questioning."
+                "If you don't know the answer, do your best without making things up. Keep the conversation engaging and inquisitive."
+                "Use chat history and context as guides but avoid repeating past responses. Provide links from the source_file metadata when relevant. Use the source context that is most relevant."
+                "Speak in a friendly and engaging manner, encouraging critical thinking and self-discovery."
+                "Use questions to lead the student to explore the topic and uncover answers."
+                "Chat History:\n{chat_history}\n\n"
+                "Context:\n{context}\n\n"
+                "Answer the student's question below by guiding them through a series of questions and insights that lead to deeper understanding. Use the context and history only if relevant, otherwise, engage in a free-flowing conversation."
+                "Foster an inquisitive mindset and help the student discover answers through dialogue."
+                "Student: {input}\n"
+                "AI Tutor:"
+            ),
+        },
+        "prompt_no_history": (
+            "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. "
+            "Provide links from the source_file metadata. Use the source context that is most relevant. "
+            "Speak in a friendly and engaging manner, like talking to a friend. Avoid sounding repetitive or robotic.\n\n"
+            "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"
+            "Student: {input}\n"
+            "AI Tutor:"
+        ),
+    },
+    "tiny_llama": {
+        "prompt_no_history": (
+            "system\n"
+            "Assistant is an intelligent chatbot designed to help students with questions regarding 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.\n"
+            "If you don't know the answer, do your best without making things up. Keep the conversation flowing naturally.\n"
+            "Provide links from the source_file metadata. Use the source context that is most relevant.\n"
+            "Speak in a friendly and engaging manner, like talking to a friend. Avoid sounding repetitive or robotic.\n"
+            "\n\n"
+            "user\n"
+            "Context:\n{context}\n\n"
+            "Question: {input}\n"
+            "\n\n"
+            "assistant"
+        ),
+        "prompt_with_history": (
+            "system\n"
+            "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 relevant. "
+            "Speak in a friendly and engaging manner, like talking to a friend. Avoid sounding repetitive or robotic.\n"
+            "\n\n"
+            "user\n"
+            "Chat History:\n{chat_history}\n\n"
+            "Context:\n{context}\n\n"
+            "Question: {input}\n"
+            "\n\n"
+            "assistant"
+        ),
+    },
+}

code/modules/dataloader/data_loader.py

@@ -63,7 +63,7 @@ class HTMLReader:
                 href = href.replace("http", "https")
 
             absolute_url = urljoin(base_url, href)
-            link['href'] = absolute_url
+            link["href"] = absolute_url
 
             resp = requests.head(absolute_url)
             if resp.status_code != 200:
@@ -84,6 +84,7 @@ class HTMLReader:
         else:
             return None
 
+
 class FileReader:
     def __init__(self, logger, kind):
         self.logger = logger
@@ -95,8 +96,9 @@ class FileReader:
         else:
             self.pdf_reader = PDFReader()
         self.web_reader = HTMLReader()
-        self.logger.info(f"Initialized FileReader with {kind} PDF reader and HTML reader")
-
+        self.logger.info(
+            f"Initialized FileReader with {kind} PDF reader and HTML reader"
+        )
 
     def extract_text_from_pdf(self, pdf_path):
         text = ""
@@ -271,11 +273,11 @@ class ChunkProcessor:
 
             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}
-            )
+
+            # Create a new dictionary for metadata in each iteration
+            metadata = addl_metadata.get(file_path, {}).copy()
+            metadata["page"] = page_num
+            metadata["source"] = file_path
             file_metadata[page_num] = metadata
 
             if self.config["vectorstore"]["db_option"] not in ["RAGatouille"]:
@@ -374,7 +376,9 @@ class ChunkProcessor:
 
 class DataLoader:
     def __init__(self, config, logger=None):
-        self.file_reader = FileReader(logger=logger, kind=config["llm_params"]["pdf_reader"])
+        self.file_reader = FileReader(
+            logger=logger, kind=config["llm_params"]["pdf_reader"]
+        )
         self.chunk_processor = ChunkProcessor(config, logger=logger)
 
     def get_chunks(self, uploaded_files, weblinks):
@@ -392,19 +396,22 @@ if __name__ == "__main__":
     with open("../code/modules/config/config.yml", "r") as f:
         config = yaml.safe_load(f)
 
-    STORAGE_DIR = os.path.join(BASE_DIR, config['vectorstore']["data_path"])
+    STORAGE_DIR = os.path.join(BASE_DIR, config["vectorstore"]["data_path"])
     uploaded_files = [
-        os.path.join(STORAGE_DIR, file) for file in os.listdir(STORAGE_DIR) if file != "urls.txt"
+        os.path.join(STORAGE_DIR, file)
+        for file in os.listdir(STORAGE_DIR)
+        if file != "urls.txt"
     ]
 
     data_loader = DataLoader(config, logger=logger)
     document_chunks, document_names, documents, document_metadata = (
         data_loader.get_chunks(
-            ["https://dl4ds.github.io/sp2024/static_files/lectures/05_loss_functions_v2.pdf"],
+            [
+                "https://dl4ds.github.io/sp2024/static_files/discussion_slides/00_discussion.pdf"
+            ],
             [],
         )
     )
 
     print(document_names[:5])
     print(len(document_chunks))
-

code/modules/vectorstore/base.py

@@ -29,5 +29,8 @@ class VectorStoreBase:
         """
         raise NotImplementedError
 
+    def __len__(self):
+        raise NotImplementedError
+
     def __str__(self):
         return self.__class__.__name__

code/modules/vectorstore/chroma.py

@@ -39,3 +39,6 @@ class ChromaVectorStore(VectorStoreBase):
 
     def as_retriever(self):
         return self.vectorstore.as_retriever()
+
+    def __len__(self):
+        return len(self.vectorstore)

code/modules/vectorstore/colbert.py

@@ -1,6 +1,67 @@
 from ragatouille import RAGPretrainedModel
 from modules.vectorstore.base import VectorStoreBase
+from langchain_core.retrievers import BaseRetriever
+from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun, Callbacks
+from langchain_core.documents import Document
+from typing import Any, List, Optional, Sequence
 import os
+import json
+
+
+class RAGatouilleLangChainRetrieverWithScore(BaseRetriever):
+    model: Any
+    kwargs: dict = {}
+
+    def _get_relevant_documents(
+        self,
+        query: str,
+        *,
+        run_manager: CallbackManagerForRetrieverRun,  # noqa
+    ) -> List[Document]:
+        """Get documents relevant to a query."""
+        docs = self.model.search(query, **self.kwargs)
+        return [
+            Document(
+                page_content=doc["content"],
+                metadata={**doc.get("document_metadata", {}), "score": doc["score"]},
+            )
+            for doc in docs
+        ]
+
+    async def _aget_relevant_documents(
+        self,
+        query: str,
+        *,
+        run_manager: CallbackManagerForRetrieverRun,  # noqa
+    ) -> List[Document]:
+        """Get documents relevant to a query."""
+        docs = self.model.search(query, **self.kwargs)
+        return [
+            Document(
+                page_content=doc["content"],
+                metadata={**doc.get("document_metadata", {}), "score": doc["score"]},
+            )
+            for doc in docs
+        ]
+
+
+class RAGPretrainedModel(RAGPretrainedModel):
+    """
+    Adding len property to RAGPretrainedModel
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._document_count = 0
+
+    def set_document_count(self, count):
+        self._document_count = count
+
+    def __len__(self):
+        return self._document_count
+
+    def as_langchain_retriever(self, **kwargs: Any) -> BaseRetriever:
+        return RAGatouilleLangChainRetrieverWithScore(model=self, kwargs=kwargs)
 
 
 class ColbertVectorStore(VectorStoreBase):
@@ -24,16 +85,28 @@ class ColbertVectorStore(VectorStoreBase):
             document_ids=document_names,
             document_metadatas=document_metadata,
         )
+        self.colbert.set_document_count(len(document_names))
 
     def load_database(self):
         path = os.path.join(
+            os.getcwd(),
             self.config["vectorstore"]["db_path"],
             "db_" + self.config["vectorstore"]["db_option"],
         )
         self.vectorstore = RAGPretrainedModel.from_index(
             f"{path}/colbert/indexes/new_idx"
         )
+
+        index_metadata = json.load(
+            open(f"{path}/colbert/indexes/new_idx/0.metadata.json")
+        )
+        num_documents = index_metadata["num_passages"]
+        self.vectorstore.set_document_count(num_documents)
+
         return self.vectorstore
 
     def as_retriever(self):
         return self.vectorstore.as_retriever()
+
+    def __len__(self):
+        return len(self.vectorstore)

code/modules/vectorstore/faiss.py

@@ -3,6 +3,13 @@ from modules.vectorstore.base import VectorStoreBase
 import os
 
 
+class FAISS(FAISS):
+    """To add length property to FAISS class"""
+
+    def __len__(self):
+        return self.index.ntotal
+
+
 class FaissVectorStore(VectorStoreBase):
     def __init__(self, config):
         self.config = config
@@ -43,3 +50,6 @@ class FaissVectorStore(VectorStoreBase):
 
     def as_retriever(self):
         return self.vectorstore.as_retriever()
+
+    def __len__(self):
+        return len(self.vectorstore)

code/modules/vectorstore/raptor.py

@@ -5,7 +5,7 @@ import os
 import numpy as np
 import pandas as pd
 import umap
-from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.prompts.chat import ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from sklearn.mixture import GaussianMixture
 from langchain_community.chat_models import ChatOpenAI
@@ -16,6 +16,13 @@ from modules.vectorstore.base import VectorStoreBase
 RANDOM_SEED = 42
 
 
+class FAISS(FAISS):
+    """To add length property to FAISS class"""
+
+    def __len__(self):
+        return self.index.ntotal
+
+
 class RAPTORVectoreStore(VectorStoreBase):
     def __init__(self, config, documents=[], text_splitter=None, embedding_model=None):
         self.documents = documents

code/modules/vectorstore/store_manager.py

@@ -3,6 +3,7 @@ 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.config.constants import RETRIEVER_HF_PATHS
 from modules.vectorstore.embedding_model_loader import EmbeddingModelLoader
 import logging
 import os
@@ -135,21 +136,33 @@ class VectorStoreManager:
             self.embedding_model = self.create_embedding_model()
         else:
             self.embedding_model = None
-        self.loaded_vector_db = self.vector_db._load_database(self.embedding_model)
+        try:
+            self.loaded_vector_db = self.vector_db._load_database(self.embedding_model)
+        except Exception as e:
+            raise ValueError(
+                f"Error loading database, check if it exists. if not run python -m modules.vectorstore.store_manager / Resteart the HF Space: {e}"
+            )
+            # print(f"Creating database")
+            # self.create_database()
+            # 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"
+            f"Time taken to load database {self.config['vectorstore']['db_option']}: {end_time - start_time} seconds"
         )
         self.logger.info("Loaded database")
         return self.loaded_vector_db
 
-    def load_from_HF(self):
+    def load_from_HF(self, HF_PATH):
         start_time = time.time()  # Start time for loading database
-        self.vector_db._load_from_HF()
+        self.vector_db._load_from_HF(HF_PATH)
         end_time = time.time()
         self.logger.info(
-            f"Time taken to load database from Hugging Face: {end_time - start_time} seconds"
+            f"Time taken to Download database {self.config['vectorstore']['db_option']} from Hugging Face: {end_time - start_time} seconds"
         )
+        self.logger.info("Downloaded database")
+
+    def __len__(self):
+        return len(self.vector_db)
 
 
 if __name__ == "__main__":
@@ -160,8 +173,18 @@ if __name__ == "__main__":
     print(config)
     print(f"Trying to create database with config: {config}")
     vector_db = VectorStoreManager(config)
-    if config["vectorstore"]["load_from_HF"] and "HF_path" in config["vectorstore"]:
-        vector_db.load_from_HF()
+    if config["vectorstore"]["load_from_HF"]:
+        if config["vectorstore"]["db_option"] in RETRIEVER_HF_PATHS:
+            vector_db.load_from_HF(
+                HF_PATH=RETRIEVER_HF_PATHS[config["vectorstore"]["db_option"]]
+            )
+        else:
+            # print(f"HF_PATH not available for {config['vectorstore']['db_option']}")
+            # print("Creating database")
+            # vector_db.create_database()
+            raise ValueError(
+                f"HF_PATH not available for {config['vectorstore']['db_option']}"
+            )
     else:
         vector_db.create_database()
     print("Created database")

code/modules/vectorstore/vectorstore.py

@@ -53,11 +53,11 @@ class VectorStore:
         else:
             return self.vectorstore.load_database(embedding_model)
 
-    def _load_from_HF(self):
+    def _load_from_HF(self, HF_PATH):
         # Download the snapshot from Hugging Face Hub
         # Note: Download goes to the cache directory
         snapshot_path = snapshot_download(
-            repo_id=self.config["vectorstore"]["HF_path"],
+            repo_id=HF_PATH,
             repo_type="dataset",
             force_download=True,
         )
@@ -86,3 +86,6 @@ class VectorStore:
 
     def _get_vectorstore(self):
         return self.vectorstore
+
+    def __len__(self):
+        return self.vectorstore.__len__()

code/public/test.css

@@ -31,3 +31,13 @@ a[href*='https://github.com/Chainlit/chainlit'] {
 .MuiAvatar-root.MuiAvatar-circular.css-v72an7 .MuiAvatar-img.css-1hy9t21 {
     display: none;
 }
+
+/* Hide the new chat button
+#new-chat-button {
+    display: none;
+} */
+
+/* Hide the open sidebar button
+#open-sidebar-button {
+    display: none;
+} */

requirements.txt

@@ -1,27 +1,25 @@
-# Automatically generated by https://github.com/damnever/pigar.
-
-aiohttp==3.9.5
-beautifulsoup4==4.12.3
-chainlit==1.1.304
-langchain==0.1.20
-langchain-community==0.0.38
-langchain-core==0.1.52
-literalai==0.0.604
-llama-parse==0.4.4
-numpy==1.26.4
-pandas==2.2.2
-pysrt==1.1.2
-python-dotenv==1.0.1
-PyYAML==6.0.1
-RAGatouille==0.0.8.post2
-requests==2.32.3
-scikit-learn==1.5.0
-torch==2.3.1
-tqdm==4.66.4
-transformers==4.41.2
-trulens_eval==0.31.0
-umap-learn==0.5.6
-trulens-eval==0.31.0
-llama-cpp-python==0.2.77
-pymupdf==1.24.5
+aiohttp
+beautifulsoup4
+chainlit
+langchain
+langchain-community
+langchain-core
+literalai
+llama-parse
+numpy
+pandas
+pysrt
+python-dotenv
+PyYAML
+RAGatouille
+requests
+scikit-learn
+torch
+tqdm
+transformers
+trulens_eval
+umap-learn
+llama-cpp-python
+pymupdf
 websockets
+langchain-openai