improvements, refactored chat

code/main.py

@@ -45,6 +45,11 @@ class Chatbot:
         """From the session `llm_settings`, create new LLMConfig and LLM objects,
         save them in session state."""
 
+        old_config = self.config.copy()  # create a copy of the previous config
+        new_config = (
+            self.config.copy()
+        )  # create the new config as a copy of the previous config
+
         llm_settings = cl.user_session.get("llm_settings", {})
         chat_profile = llm_settings.get("chat_model")
         retriever_method = llm_settings.get("retriever_method")
@@ -54,16 +59,18 @@ class Chatbot:
 
         chain = cl.user_session.get("chain")
         memory = chain.memory
-        self.config["vectorstore"][
+        new_config["vectorstore"][
             "db_option"
         ] = retriever_method  # update the retriever method in the config
-        memory.k = memory_window  # set the memory window
+        new_config["llm_params"][
+            "memory_window"
+        ] = memory_window  # update the memory window in the config
 
-        self.llm_tutor = LLMTutor(self.config)
+        self.llm_tutor.update_llm(new_config)
         self.chain = self.llm_tutor.qa_bot(memory=memory)
 
         tags = [chat_profile, self.config["vectorstore"]["db_option"]]
-        self.chat_processor = ChatProcessor(self.config, tags=tags)
+        self.chat_processor = ChatProcessor(self.llm_tutor, tags=tags)
 
         cl.user_session.set("chain", self.chain)
         cl.user_session.set("llm_tutor", self.llm_tutor)
@@ -103,6 +110,11 @@ class Chatbot:
                 cl.input_widget.Switch(
                     id="view_sources", label="View Sources", initial=False
                 ),
+                # cl.input_widget.TextInput(
+                #     id="vectorstore",
+                #     label="temp",
+                #     initial="None",
+                # ),
             ]
         ).send()  # type: ignore
 
@@ -156,14 +168,14 @@ class Chatbot:
 
     async def chat_profile(self):
         return [
-            # 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="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**.",
@@ -181,37 +193,44 @@ class Chatbot:
         if chat_profile:
             self._configure_llm(chat_profile)
 
-        self.llm_tutor = LLMTutor(self.config)
+        self.llm_tutor = LLMTutor(
+            self.config, user={"user_id": "abc123", "session_id": "789"}
+        )
         self.chain = self.llm_tutor.qa_bot()
         tags = [chat_profile, self.config["vectorstore"]["db_option"]]
-        self.chat_processor = ChatProcessor(self.config, tags=tags)
+        self.chat_processor = ChatProcessor(self.llm_tutor, tags=tags)
 
         cl.user_session.set("llm_tutor", self.llm_tutor)
         cl.user_session.set("chain", self.chain)
-        cl.user_session.set("counter", 0)
+        cl.user_session.set("counter", 20)
         cl.user_session.set("chat_processor", self.chat_processor)
 
     async def on_chat_end(self):
         await cl.Message(content="Sorry, I have to go now. Goodbye!").send()
 
     async def main(self, message):
-        user = cl.user_session.get("user")
         chain = cl.user_session.get("chain")
         counter = cl.user_session.get("counter")
-        llm_settings = cl.user_session.get("llm_settings")
+        llm_settings = cl.user_session.get("llm_settings", {})
+        view_sources = llm_settings.get("view_sources", False)
+
+        print("HERE")
+        print(llm_settings)
+        print(view_sources)
+        print("\n\n")
 
         counter += 1
         cl.user_session.set("counter", counter)
 
-        cb = cl.AsyncLangchainCallbackHandler()  # TODO: fix streaming here
-        cb.answer_reached = True
-
         processor = cl.user_session.get("chat_processor")
-        res = await processor.rag(message.content, chain, cb)
+        res = await processor.rag(message.content, chain)
+
+        print(res)
+
         answer = res.get("answer", res.get("result"))
 
         answer_with_sources, source_elements, sources_dict = get_sources(
-            res, answer, view_sources=llm_settings.get("view_sources")
+            res, answer, view_sources=view_sources
         )
         processor._process(message.content, answer, sources_dict)
 

code/modules/chat/helpers.py

@@ -1,13 +1,14 @@
 from modules.config.constants import *
 import chainlit as cl
 from langchain_core.prompts import PromptTemplate
+from langchain_core.prompts import ChatPromptTemplate
 
 
 def get_sources(res, answer, 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")
@@ -43,64 +44,77 @@ def get_sources(res, answer, view_sources=False):
     if view_sources:
 
         # Then, display the sources
-        full_answer += "\n\n**Sources:**\n"
-        for idx, (url_name, source_data) in enumerate(source_dict.items()):
-            full_answer += f"\nSource {idx + 1} (Score: {source_data['score']}): {source_data['url']}\n"
-
-            name = f"Source {idx + 1} Text\n"
-            full_answer += name
-            source_elements.append(
-                cl.Text(name=name, content=source_data["text"], display="side")
-            )
+        # 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"
 
-            # 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"
+                name = f"Source {idx + 1} Text\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",
+                source_elements.append(
+                    cl.Text(name=name, content=source_data["text"], display="side")
+                )
+
+                # Add a PDF element if the source is a PDF file
+                if source_data["url"].lower().endswith(".pdf"):
+                    name = f"Source {idx + 1} PDF\n"
+                    full_answer += name
+                    pdf_url = f"{source_data['url']}#page={source_data['page']+1}"
+                    source_elements.append(
+                        cl.Pdf(name=name, url=pdf_url, display="side")
+                    )
+
+            full_answer += "\n**Metadata:**\n"
+            for idx, (url_name, source_data) in enumerate(source_dict.items()):
+                full_answer += f"\nSource {idx + 1} Metadata:\n"
+                source_elements.append(
+                    cl.Text(
+                        name=f"Source {idx + 1} Metadata",
+                        content=f"Source: {source_data['url']}\n"
+                        f"Page: {source_data['page']}\n"
+                        f"Type: {source_data['source_type']}\n"
+                        f"Date: {source_data['date']}\n"
+                        f"TL;DR: {source_data['lecture_tldr']}\n"
+                        f"Lecture Recording: {source_data['lecture_recording']}\n"
+                        f"Suggested Readings: {source_data['suggested_readings']}\n",
+                        display="side",
+                    )
                 )
-            )
 
     return full_answer, source_elements, source_dict
 
 
-def get_prompt(config):
-    if config["llm_params"]["use_history"]:
-        if config["llm_params"]["llm_loader"] == "local_llm":
-            custom_prompt_template = tinyllama_prompt_template_with_history
-        elif config["llm_params"]["llm_loader"] == "openai":
-            custom_prompt_template = openai_prompt_template_with_history
-        # else:
-        #     custom_prompt_template = tinyllama_prompt_template_with_history # default
-        prompt = PromptTemplate(
-            template=custom_prompt_template,
-            input_variables=["context", "chat_history", "question"],
-        )
-    else:
-        if config["llm_params"]["llm_loader"] == "local_llm":
-            custom_prompt_template = tinyllama_prompt_template
-        elif config["llm_params"]["llm_loader"] == "openai":
-            custom_prompt_template = openai_prompt_template
-        # else:
-        #     custom_prompt_template = tinyllama_prompt_template
-        prompt = PromptTemplate(
-            template=custom_prompt_template,
-            input_variables=["context", "question"],
+def get_prompt(config, prompt_type):
+    llm_params = config["llm_params"]
+    llm_loader = llm_params["llm_loader"]
+    use_history = llm_params["use_history"]
+
+    if prompt_type == "qa":
+        if llm_loader == "openai":
+            return (
+                OPENAI_PROMPT_WITH_HISTORY if use_history else OPENAI_PROMPT_NO_HISTORY
+            )
+        elif (
+            llm_loader == "local_llm"
+            and llm_params.get("local_llm_params") == "tiny-llama"
+        ):
+            return (
+                TINYLLAMA_PROMPT_TEMPLATE_WITH_HISTORY
+                if use_history
+                else TINYLLAMA_PROMPT_TEMPLATE_NO_HISTORY
+            )
+    elif prompt_type == "rephrase":
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                ("system", OPENAI_REPHRASE_PROMPT),
+                ("human", "{question}, {chat_history}"),
+            ]
         )
-    return prompt
+        return OPENAI_REPHRASE_PROMPT
+
+    return None

code/modules/chat/langchain/langchain_rag.py

@@ -0,0 +1,137 @@
+from langchain_core.prompts import ChatPromptTemplate
+
+from modules.chat.langchain.utils import *
+
+
+class CustomConversationalRetrievalChain:
+    def __init__(self, llm, memory, retriever, qa_prompt: str, rephrase_prompt: str):
+        """
+        Initialize the CustomConversationalRetrievalChain 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 = 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_messages(
+            [
+                ("system", contextualize_q_system_prompt),
+                MessagesPlaceholder("chat_history"),
+                ("human", "{input}"),
+            ]
+        )
+
+        # 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_messages(
+            [
+                ("system", qa_system_prompt),
+                MessagesPlaceholder("chat_history"),
+                ("human", "{input}"),
+            ]
+        )
+
+        # 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,
+                ),
+            ],
+        )
+
+    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()
+        return self.store[(user_id, conversation_id)]
+
+    def invoke(self, user_query, config):
+        """
+        Invoke the chain.
+
+        Args:
+            kwargs: The input variables.
+
+        Returns:
+            dict: The output variables.
+        """
+        print(user_query, config)
+        return self.rag_chain.invoke(user_query, config)

code/modules/chat/langchain/utils.py

@@ -0,0 +1,197 @@
+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
+
+
+class CustomRunnableWithHistory(RunnableWithMessageHistory):
+    def _enter_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 = hist.messages.copy()
+
+        if not self.history_messages_key:
+            # return all messages
+            messages += self._get_input_messages(input)
+
+        # 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"] :]
+        return messages
+
+
+def _get_chat_history(chat_history: List[CHAT_TURN_TYPE], n: int = None) -> str:
+    """
+    Convert chat history to a formatted string.
+
+    Args:
+        chat_history (List[CHAT_TURN_TYPE]): The chat history.
+
+    Returns:
+        str: The formatted chat history.
+    """
+    _ROLE_MAP = {"human": "Student: ", "ai": "AI Tutor: "}
+    buffer = ""
+    if n is not None:
+        # Calculate the number of turns to take (2 turns per pair)
+        turns_to_take = n * 2
+        chat_history = chat_history[-turns_to_take:]
+    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
+
+
+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 get_last_n_conversations(self, n: int) -> "InMemoryHistory":
+        """Return a new InMemoryHistory object with the last n conversations from the message history.
+
+        Args:
+            n (int): The number of last conversations to return. If 0, return an empty history.
+
+        Returns:
+            InMemoryHistory: A new InMemoryHistory object containing the last n conversations.
+        """
+        if n == 0:
+            return InMemoryHistory()
+        # Each conversation consists of a pair of messages (human + AI)
+        num_messages = n * 2
+        last_messages = self.messages[-num_messages:]
+        return InMemoryHistory(messages=last_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

code/modules/chat/llm_tutor.py

@@ -1,216 +1,140 @@
-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 CustomConversationalRetrievalChain
 
 
 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)
         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, 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(self.config, new_config)
+        self.config = new_config
+
+        if "chat_model" in changes:
+            self.llm = self.load_llm()  # Reinitialize LLM if chat_model changes
 
-        return prompt
+        if "vectorstore" in changes:
+            self.vector_db = VectorStoreManager(
+                self.config, logger=self.logger
+            )  # 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, memory=None):
+    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 = {}
+        for key in new_config:
+            if old_config.get(key) != new_config[key]:
+                changes[key] = (old_config.get(key), new_config[key])
+        return changes
+
+    def retrieval_qa_chain(self, llm, qa_prompt, rephrase_prompt, db, memory=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"]:
-            if memory is None:
-                memory = ConversationBufferWindowMemory(
-                    k=self.config["llm_params"]["memory_window"],
-                    memory_key="chat_history",
-                    return_messages=True,
-                    output_key="answer",
-                )
-            qa_chain = CustomConversationalRetrievalChain.from_llm(
+            qa_chain = CustomConversationalRetrievalChain(
                 llm=llm,
-                chain_type="stuff",
-                retriever=retriever,
-                return_source_documents=True,
                 memory=memory,
-                combine_docs_chain_kwargs={"prompt": prompt},
-                response_if_no_docs_found="No context found",
-            )
-        else:
-            qa_chain = RetrievalQA.from_chain_type(
-                llm=llm,
-                chain_type="stuff",
                 retriever=retriever,
-                return_source_documents=True,
-                chain_type_kwargs={"prompt": prompt},
+                qa_prompt=qa_prompt,
+                rephrase_prompt=rephrase_prompt,
             )
         return 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, memory=None):
+    def qa_bot(self, memory=None, qa_prompt=None, rephrase_prompt=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.
+        """
+        if qa_prompt is None:
+            qa_prompt = get_prompt(self.config, "qa")
+        if rephrase_prompt is None:
+            rephrase_prompt = get_prompt(self.config, "rephrase")
         db = self.vector_db.load_database()
         # sanity check to see if there are any documents in the database
         if len(db) == 0:
             raise ValueError(
                 "No documents in the database. Populate the database first."
             )
-        qa_prompt = self.set_custom_prompt()
-        qa = self.retrieval_qa_chain(
-            self.llm, qa_prompt, db, memory
-        )  # TODO:  PROMPT is overwritten in CustomConversationalRetrievalChain
+        qa = self.retrieval_qa_chain(self.llm, qa_prompt, rephrase_prompt, db, memory)
 
         return qa
 
-    # output function
     def final_result(query):
+        """
+        Get the final result for a given query.
+
+        Args:
+            query (str): The query string.
+
+        Returns:
+            str: The response string.
+        """
         qa_result = qa_bot()
         response = qa_result({"query": query})
         return response

code/modules/chat_processor/chat_processor.py

@@ -2,13 +2,25 @@ 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
+    def __init__(self, llm_tutor, tags=None):
+        self.llm_tutor = llm_tutor
+        self.config = self.llm_tutor.config
+        self.chat_processor_type = self.config["chat_logging"]["platform"]
+        self.logging = self.config["chat_logging"]["log_chat"]
+        self.user = self.llm_tutor.user
+        if tags is None:
+            self.tags = self._create_tags()
+        else:
+            self.tags = tags
         if self.logging:
             self._init_processor()
 
+    def _create_tags(self):
+        tags = []
+        tags.append(self.config["vectorstore"]["db_option"])
+        tags.append(self.config["llm_params"]["chat_profile"])
+        return tags
+
     def _init_processor(self):
         if self.chat_processor_type == "literalai":
             self.processor = LiteralaiChatProcessor(self.tags)
@@ -23,8 +35,21 @@ class ChatProcessor:
         else:
             pass
 
-    async def rag(self, user_query: str, chain, cb):
+    async def rag(self, user_query: str, chain):
+        user_query_dict = {"input": user_query}
+        # Define the base configuration
+        config = {
+            "configurable": {
+                "user_id": self.user["user_id"],
+                "conversation_id": self.user["session_id"],
+                "memory_window": self.llm_tutor.config["llm_params"]["memory_window"],
+            }
+        }
+
+        # Process the user query using the appropriate method
         if self.logging:
-            return await self.processor.rag(user_query, chain, cb)
+            return await self.processor.rag(
+                user_query=user_query_dict, config=config, chain=chain
+            )
         else:
-            return await chain.acall(user_query, callbacks=[cb])
+            return chain.invoke(user_query=user_query_dict, config=config)

code/modules/chat_processor/literal_ai.py

@@ -27,11 +27,11 @@ class LiteralaiChatProcessor(ChatProcessorBase):
                 name="AI_Tutor",
             )
 
-    async def rag(self, user_query: str, chain, cb):
+    async def rag(self, user_query: dict, config: dict, chain):
         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.input = {"question": user_query["input"]}
+            res = await chain.invoke(user_query, config)
             step.output = res
         return res

code/modules/config/constants.py

@@ -13,70 +13,69 @@ opening_message = f"Hey, What Can I Help You With?\n\nYou can me ask me question
 
 # Prompt Templates
 
-openai_prompt_template = """Use the following pieces of information to answer the user's question.
-If you don't know the answer, just say that you don't know.
-
-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.
-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 breif and concise answer to the question. Use the history to answer the question if you can.
-
-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 breif and concise answer to the question.
-
-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
-"""
-
-
+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}'"
+)
+
+OPENAI_PROMPT_WITH_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. "
+    "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"
+    "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:"
+)
+
+OPENAAI_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:"
+)
+
+
+TINYLLAMA_PROMPT_TEMPLATE_NO_HISTORY = (
+    "<|im_start|>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"
+    "<|im_end|>\n\n"
+    "<|im_start|>user\n"
+    "Context:\n{context}\n\n"
+    "Question: {input}\n"
+    "<|im_end|>\n\n"
+    "<|im_start|>assistant"
+)
+
+TINYLLAMA_PROMPT_TEMPLATE_WITH_HISTORY = (
+    "<|im_start|>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"
+    "<|im_end|>\n\n"
+    "<|im_start|>user\n"
+    "Chat History:\n{chat_history}\n\n"
+    "Context:\n{context}\n\n"
+    "Question: {input}\n"
+    "<|im_end|>\n\n"
+    "<|im_start|>assistant"
+)
 # Model Paths
 
 LLAMA_PATH = "../storage/models/tinyllama-1.1b-chat-v1.0.Q5_K_M.gguf"
-MISTRAL_PATH = "storage/models/mistral-7b-v0.1.Q4_K_M.gguf"