saving commit - needs cleanup

code/main.py

@@ -1,5 +1,5 @@
 import chainlit.data as cl_data
-
+import asyncio
 from modules.config.constants import (
     LLAMA_PATH,
     LITERAL_API_KEY_LOGGING,
@@ -13,10 +13,15 @@ import os
 from typing import Any, Dict, no_type_check
 import chainlit as cl
 from modules.chat.llm_tutor import LLMTutor
-from modules.chat.helpers import get_sources
+from modules.chat.helpers import (
+    get_sources,
+    get_history_chat_resume,
+    get_history_setup_llm,
+)
 import copy
 from typing import Optional
 from chainlit.types import ThreadDict
+import time
 
 USER_TIMEOUT = 60_000
 SYSTEM = "System 🖥️"
@@ -25,19 +30,30 @@ AGENT = "Agent <>"
 YOU = "You 😃"
 ERROR = "Error 🚫"
 
+with open("modules/config/config.yml", "r") as f:
+    config = yaml.safe_load(f)
 
-cl_data._data_layer = CustomLiteralDataLayer(
-    api_key=LITERAL_API_KEY_LOGGING, server=LITERAL_API_URL
-)
+
+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):
+    def __init__(self, config):
         """
         Initialize the Chatbot class.
         """
-        self.config = self._load_config()
-        self.literal_client = cl_data._data_layer.client
+        self.config = config
 
     def _load_config(self):
         """
@@ -51,6 +67,8 @@ class Chatbot:
         """
         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 = (
             llm_settings.get("chat_model"),
@@ -60,7 +78,6 @@ class Chatbot:
         )
 
         chain = cl.user_session.get("chain")
-        print(list(chain.store.values()))
         memory_list = cl.user_session.get(
             "memory",
             (
@@ -69,35 +86,7 @@ class Chatbot:
                 else []
             ),
         )
-        conversation_list = []
-        for message in memory_list:
-            # Convert to dictionary if possible
-            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)
-            )
-
-            # 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)
-            )
-
-            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")
+        conversation_list = get_history_setup_llm(memory_list)
         print("\n\n\n")
         print("history at setup_llm", conversation_list)
         print("\n\n\n")
@@ -111,13 +100,18 @@ class Chatbot:
         self.llm_tutor.update_llm(
             old_config, self.config
         )  # update only attributes that are changed
-        self.chain = self.llm_tutor.qa_bot(memory=conversation_list)
+        self.chain = self.llm_tutor.qa_bot(
+            memory=conversation_list,
+            callbacks=[cl.LangchainCallbackHandler()] if cl_data._data_layer 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]):
         """
@@ -176,7 +170,7 @@ class Chatbot:
                 cl.input_widget.Select(
                     id="llm_style",
                     label="Type of Conversation (Default Normal)",
-                    values=["Normal", "ELI5", "Socratic"],
+                    values=["Normal", "ELI5"],
                     initial_index=0,
                 ),
             ]
@@ -268,6 +262,8 @@ class Chatbot:
         and display and load previous conversation if chat logging is enabled.
         """
 
+        start_time = time.time()
+
         await self.make_llm_settings_widgets(self.config)
         user = cl.user_session.get("user")
         self.user = {
@@ -280,10 +276,20 @@ class Chatbot:
 
         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)
+
+        print(cl.LangchainCallbackHandler())
+        print(cl_data._data_layer)
+        self.chain = self.llm_tutor.qa_bot(
+            memory=memory,
+            callbacks=[cl.LangchainCallbackHandler()] if cl_data._data_layer else None,
+        )
+        self.question_generator = self.llm_tutor.question_generator
+        print(self.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.
@@ -314,6 +320,8 @@ class Chatbot:
             message: The incoming chat message.
         """
 
+        start_time = time.time()
+
         chain = cl.user_session.get("chain")
 
         print("\n\n\n")
@@ -348,64 +356,72 @@ class Chatbot:
             res = chain.stream(user_query=user_query_dict, config=chain_config)
             res = await self.stream_response(res)
         else:
-            res = await chain.invoke(user_query=user_query_dict, config=chain_config)
+            res = await chain.invoke(
+                user_query=user_query_dict,
+                config=chain_config,
+            )
 
         answer = res.get("answer", res.get("result"))
 
-        with cl_data._data_layer.client.step(
-            type="retrieval",
-            name="RAG",
-            thread_id=cl.context.session.thread_id,
-            # tags=self.tags,
-        ) as step:
-            step.input = {"question": user_query_dict["input"]}
-            step.output = {
-                "chat_history": res.get("chat_history"),
-                "context": res.get("context"),
-                "answer": answer,
-                "rephrase_prompt": res.get("rephrase_prompt"),
-                "qa_prompt": res.get("qa_prompt"),
-            }
-            step.metadata = self.config
+        if cl_data._data_layer is not None:
+            with cl_data._data_layer.client.step(
+                type="run",
+                name="step_info",
+                thread_id=cl.context.session.thread_id,
+                # tags=self.tags,
+            ) as step:
+
+                step.input = {"question": user_query_dict["input"]}
+
+                step.output = {
+                    "chat_history": res.get("chat_history"),
+                    "context": res.get("context"),
+                    "answer": answer,
+                    "rephrase_prompt": res.get("rephrase_prompt"),
+                    "qa_prompt": res.get("qa_prompt"),
+                }
+                step.metadata = self.config
 
         answer_with_sources, source_elements, sources_dict = get_sources(
             res, answer, stream=stream, view_sources=view_sources
         )
 
+        print("Time taken to process the message: ", time.time() - start_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"),
+        )
+
+        print("\n\n\n")
+        print("Questions: ", list_of_questions)
+        print("\n\n\n")
+
+        actions = []
+        for question in list_of_questions:
+
+            actions.append(
+                cl.Action(
+                    name="follow up question",
+                    value="example_value",
+                    description=question,
+                    label=question,
+                )
+            )
+
         await cl.Message(
-            content=answer_with_sources, elements=source_elements, author=LLM
+            content=answer_with_sources,
+            elements=source_elements,
+            author=LLM,
+            actions=actions,
         ).send()
 
     async def on_chat_resume(self, thread: ThreadDict):
         steps = thread["steps"]
-        # conversation_pairs = []
-        conversation_list = []
-
-        user_message = None
         k = self.config["llm_params"]["memory_window"]
-        count = 0
-
-        print(steps)
-
-        for step in reversed(steps):
-            print(step["type"])
-            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]
+        conversation_list = get_history_chat_resume(steps, k, SYSTEM, LLM)
 
         print("\n\n\n")
         print("history at on_chat_resume", conversation_list)
@@ -423,11 +439,30 @@ class Chatbot:
     ) -> Optional[cl.User]:
         return default_user
 
+    async def on_action(self, action: cl.Action):
+        print("Action Callback")
+        print(action)
+        # main(message=action.description)\
+        message = await cl.Message(content=action.description, author=YOU).send()
+        await self.main(message)
+
+
+chatbot = Chatbot(config=config)
+
+
+async def start():
+    print("Setting up data layer...")
+    cl_data._data_layer = await setup_data_layer()
+    print("Data layer set up.")
+    print(cl_data._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_action)
+
 
-chatbot = Chatbot()
-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)
+asyncio.run(start())

code/modules/chat/helpers.py

@@ -116,3 +116,58 @@ def get_prompt(config, prompt_type):
                 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):
+        print(step["type"])
+        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)
+        )
+
+        # 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)
+        )
+
+        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

code/modules/chat/langchain/langchain_rag.py

@@ -9,11 +9,21 @@ from langchain.memory import (
     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
+        self,
+        llm,
+        memory,
+        retriever,
+        qa_prompt: str,
+        rephrase_prompt: str,
+        config: dict,
+        callbacks=None,
     ):
         """
         Initialize the Langchain_RAG class.
@@ -77,9 +87,29 @@ class Langchain_RAG_V1(BaseRAG):
         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
+        self,
+        llm,
+        memory,
+        retriever,
+        qa_prompt: str,
+        rephrase_prompt: str,
+        config: dict,
+        callbacks=None,
     ):
         """
         Initialize the Langchain_RAG class.
@@ -171,6 +201,9 @@ class Langchain_RAG_V2(BaseRAG):
             ],
         )
 
+        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:
@@ -192,7 +225,7 @@ class Langchain_RAG_V2(BaseRAG):
             )  # 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):
+    async def invoke(self, user_query, config, **kwargs):
         """
         Invoke the chain.
 
@@ -202,7 +235,7 @@ class Langchain_RAG_V2(BaseRAG):
         Returns:
             dict: The output variables.
         """
-        res = await self.rag_chain.ainvoke(user_query, config)
+        res = await self.rag_chain.ainvoke(user_query, config, **kwargs)
         res["rephrase_prompt"] = self.rephrase_prompt
         res["qa_prompt"] = self.qa_prompt
         return res

code/modules/chat/langchain/utils.py

@@ -311,3 +311,39 @@ def create_retrieval_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

@@ -2,7 +2,11 @@ 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 modules.chat.langchain.langchain_rag import Langchain_RAG_V1, Langchain_RAG_V2
+from modules.chat.langchain.langchain_rag import (
+    Langchain_RAG_V1,
+    Langchain_RAG_V2,
+    QuestionGenerator,
+)
 
 
 class LLMTutor:
@@ -86,7 +90,9 @@ class LLMTutor:
         compare_dicts(old_config, new_config)
         return changes
 
-    def retrieval_qa_chain(self, llm, qa_prompt, rephrase_prompt, db, memory=None):
+    def retrieval_qa_chain(
+        self, llm, qa_prompt, rephrase_prompt, db, memory=None, callbacks=None
+    ):
         """
         Create a Retrieval QA Chain.
 
@@ -110,7 +116,10 @@ class LLMTutor:
                 qa_prompt=qa_prompt,
                 rephrase_prompt=rephrase_prompt,
                 config=self.config,
+                callbacks=callbacks,
             )
+
+            self.question_generator = QuestionGenerator()
         else:
             raise ValueError(
                 f"Invalid LLM Architecture: {self.config['llm_params']['llm_arch']}"
@@ -128,7 +137,7 @@ class LLMTutor:
         llm = chat_model_loader.load_chat_model()
         return llm
 
-    def qa_bot(self, memory=None):
+    def qa_bot(self, memory=None, callbacks=None):
         """
         Create a QA bot instance.
 
@@ -147,7 +156,12 @@ class LLMTutor:
             )
 
         qa = self.retrieval_qa_chain(
-            self.llm, self.qa_prompt, self.rephrase_prompt, self.vector_db, memory
+            self.llm,
+            self.qa_prompt,
+            self.rephrase_prompt,
+            self.vector_db,
+            memory,
+            callbacks=callbacks,
         )
 
         return qa

code/modules/config/config.yml

@@ -28,7 +28,7 @@ llm_params:
   llm_arch: 'langchain' # [langchain, langgraph_agentic]
   use_history: True # bool
   memory_window: 3 # int
-  llm_style: 'Normal' # str [Normal, ELI5, Socratic]
+  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:
     temperature: 0.7 # float
@@ -39,6 +39,7 @@ llm_params:
 chat_logging:
   log_chat: False # bool
   platform: 'literalai'
+  callbacks: True # bool
 
 splitter_options:
   use_splitter: True # bool

code/modules/config/prompts.py

@@ -24,31 +24,28 @@ prompts = {
                 "AI Tutor:"
             ),
             "eli5": (
-                "You are an AI Tutor for the course DS598, taught by Prof. Thomas Gardos. Answer the user's question using the provided context in the simplest way possible, as if you are explaining to a 5-year-old. Only use the context if it helps make things clearer. The context is ordered by relevance. "
-                "If you don't know the answer, do your best without making things up. Keep the conversation simple and easy to understand. "
-                "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 curious child. Avoid complex terms.\n\n"
+                "You are an AI Tutor for the course DS598, taught by Prof. Thomas Gardos. Answer the user's question in the simplest way possible, like explaining to someone new to the topic. Use the provided context to help clarify your explanation."
+                "If you don't know the answer, do your best without making things up. Keep the conversation straightforward and easy to follow."
+                "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, like talking to someone who is curious and eager to learn. Avoid using complex terms and jargon."
+                "Use examples wherever possible to aid in understanding."
                 "Chat History:\n{chat_history}\n\n"
                 "Context:\n{context}\n\n"
-                "Answer the student's question below in a friendly, simple, and engaging manner. Use the context and history only if relevant, otherwise, engage in a free-flowing conversation.\n"
-                "Give a long very detailed narrative explanation. Use examples wherever you can to aid in the explanation. Remember, explain it as if you are talking to a 5-year-old, so construct a long narrative that builds up to the answer.\n"
-                "5-year-old Student: {input}\n"
+                "Answer the student's question below in a friendly, simple, and engaging manner. Use the context and history only if relevant, otherwise, engage in a free-flowing conversation."
+                "Give a clear and detailed explanation with examples to make it easier to understand."
+                "Student: {input}\n"
                 "AI Tutor:"
             ),
             "socratic": (
-                "You are an AI Tutor for the course DS598, taught by Prof. Thomas Gardos. Your goal is to guide the student towards understanding using the Socratic method. Ask thought-provoking questions to encourage critical thinking and self-discovery. Use the provided context only when relevant. The context is ordered by relevance.\n\n"
-                "Guidelines for the Socratic approach:\n"
-                "Guidelines:"
-                "1. Begin with a concise, direct answer to the student's question."
-                "2. Follow up with 1-2 thought-provoking questions to encourage critical thinking."
-                "3. Provide additional explanations or context if necessary to move the conversation forward."
-                "4. End with an open-ended question that invites further exploration."
-                "Based on the chat history determine which guideline to follow., and answer accordingly\n\n"
-                "If the student is stuck, offer gentle hints or break down the concept into simpler parts. Maintain a friendly, engaging tone throughout the conversation.\n\n"
-                "Use chat history and context as guides, but avoid repeating past responses. Provide links from the source_file metadata when appropriate. Use the most relevant source context.\n\n"
+                "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"
-                "Engage with the student's question below using the Socratic method. Ask probing questions to guide their thinking and encourage deeper understanding. Only provide direct answers if absolutely necessary.\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:"
             ),