updates

code/main.py

@@ -1,14 +1,22 @@
+import chainlit.data as cl_data
+
+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.chat_processor.chat_processor import ChatProcessor
-from modules.config.constants import LLAMA_PATH
 from modules.chat.helpers import get_sources
 import copy
 from typing import Optional
+from chainlit.types import ThreadDict
 
 USER_TIMEOUT = 60_000
 SYSTEM = "System 🖥️"
@@ -18,12 +26,18 @@ YOU = "You 😃"
 ERROR = "Error 🚫"
 
 
+cl_data._data_layer = CustomLiteralDataLayer(
+    api_key=LITERAL_API_KEY_LOGGING, server=LITERAL_API_URL
+)
+
+
 class Chatbot:
     def __init__(self):
         """
         Initialize the Chatbot class.
         """
         self.config = self._load_config()
+        self.literal_client = cl_data._data_layer.client
 
     def _load_config(self):
         """
@@ -60,11 +74,9 @@ class Chatbot:
         self.chain = self.llm_tutor.qa_bot(memory=memory)
 
         tags = [chat_profile, self.config["vectorstore"]["db_option"]]
-        self.chat_processor.config = self.config
 
         cl.user_session.set("chain", self.chain)
         cl.user_session.set("llm_tutor", self.llm_tutor)
-        cl.user_session.set("chat_processor", self.chat_processor)
 
     @no_type_check
     async def update_llm(self, new_settings: Dict[str, Any]):
@@ -91,14 +103,21 @@ class Chatbot:
                 cl.input_widget.Select(
                     id="chat_model",
                     label="Model Name (Default GPT-3)",
-                    values=["local_llm", "gpt-3.5-turbo-1106", "gpt-4"],
-                    initial_index=["local_llm", "gpt-3.5-turbo-1106", "gpt-4"].index(config["llm_params"]["llm_loader"]),
+                    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"])
+                    initial_index=["FAISS", "Chroma", "RAGatouille", "RAPTOR"].index(
+                        config["vectorstore"]["db_option"]
+                    ),
                 ),
                 cl.input_widget.Slider(
                     id="memory_window",
@@ -112,7 +131,7 @@ class Chatbot:
                     id="view_sources", label="View Sources", initial=False
                 ),
                 cl.input_widget.Switch(
-                    id="stream_response", label="Stream response", initial=True
+                    id="stream_response", label="Stream response", initial=False
                 ),
                 cl.input_widget.Select(
                     id="llm_style",
@@ -158,28 +177,37 @@ class Chatbot:
         """
         Set starter messages for the chatbot.
         """
-        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",
-            ),
-        ]
+        # 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):
         """
@@ -194,44 +222,25 @@ class Chatbot:
         rename_dict = {"Chatbot": "AI Tutor"}
         return rename_dict.get(orig_author, orig_author)
 
-    async def start(self):
+    async def start(self, thread=None, memory=[]):
         """
         Start the chatbot, initialize settings widgets,
         and display and load previous conversation if chat logging is enabled.
         """
-        await cl.Message(content="Welcome back! Setting up your session...").send()
 
         await self.make_llm_settings_widgets(self.config)
         user = cl.user_session.get("user")
         self.user = {
             "user_id": user.identifier,
-            "session_id": "1234",
+            "session_id": cl.context.session.thread_id,
         }
+        print(self.user)
+
         cl.user_session.set("user", self.user)
-        self.chat_processor = ChatProcessor(self.config, self.user)
         self.llm_tutor = LLMTutor(self.config, user=self.user)
-        if self.config["chat_logging"]["log_chat"]:
-            # get previous conversation of the user
-            memory = self.chat_processor.processor.prev_conv
-            if len(self.chat_processor.processor.prev_conv) > 0:
-                for idx, conv in enumerate(self.chat_processor.processor.prev_conv):
-                    await cl.Message(
-                        author="User", content=conv[0], type="user_message"
-                    ).send()
-                    await cl.Message(author="AI Tutor", content=conv[1]).send()
-        else:
-            memory = []
         self.chain = self.llm_tutor.qa_bot(memory=memory)
         cl.user_session.set("llm_tutor", self.llm_tutor)
         cl.user_session.set("chain", self.chain)
-        cl.user_session.set("chat_processor", self.chat_processor)
-
-    async def on_chat_end(self):
-        """
-        Handle the end of the chat session by sending a goodbye message.
-        # TODO: Not used as of now - useful when the implementation for the conversation limiting is implemented
-        """
-        await cl.Message(content="Sorry, I have to go now. Goodbye!").send()
 
     async def stream_response(self, response):
         """
@@ -245,8 +254,8 @@ class Chatbot:
 
         output = {}
         for chunk in response:
-            if 'answer' in chunk:
-                await msg.stream_token(chunk['answer'])
+            if "answer" in chunk:
+                await msg.stream_token(chunk["answer"])
 
             for key in chunk:
                 if key not in output:
@@ -262,39 +271,88 @@ class Chatbot:
         Args:
             message: The incoming chat message.
         """
+
         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", True)) or (not self.config["llm_params"]["stream"])
-
-        processor = cl.user_session.get("chat_processor")
-        res = await processor.rag(message.content, chain, stream)
+        stream = (llm_settings.get("stream_response", True)) or (
+            not self.config["llm_params"]["stream"]
+        )
+        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:
+            res = 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
+
         answer_with_sources, source_elements, sources_dict = get_sources(
             res, answer, stream=stream, view_sources=view_sources
         )
-        processor._process(message.content, answer, sources_dict)
 
         await cl.Message(content=answer_with_sources, elements=source_elements).send()
 
+    async def on_chat_resume(self, thread: ThreadDict):
+        steps = thread["steps"]
+        conversation_pairs = []
+
+        user_message = None
+        k = self.config["llm_params"]["memory_window"]
+        count = 0
+
+        for step in steps:
+            if step["type"] == "user_message":
+                user_message = step["output"]
+            elif step["type"] == "assistant_message" and user_message is not None:
+                assistant_message = step["output"]
+                conversation_pairs.append((user_message, assistant_message))
+                user_message = None
+                count += 1
+                if count >= k:
+                    break
+
+        await self.start(thread, memory=conversation_pairs)
+
     @cl.oauth_callback
     def auth_callback(
-            provider_id: str,
-            token: str,
-            raw_user_data: Dict[str, str],
-            default_user: cl.User,
+        provider_id: str,
+        token: str,
+        raw_user_data: Dict[str, str],
+        default_user: cl.User,
     ) -> Optional[cl.User]:
         return default_user
 
+
 chatbot = Chatbot()
 cl.set_starters(chatbot.set_starters)
 cl.author_rename(chatbot.rename)
 cl.on_chat_start(chatbot.start)
-cl.on_chat_end(chatbot.on_chat_end)
+cl.on_chat_resume(chatbot.on_chat_resume)
 cl.on_message(chatbot.main)
 cl.on_settings_update(chatbot.update_llm)

code/modules/chat/chat_model_loader.py

@@ -16,7 +16,11 @@ class ChatModelLoader:
         self.huggingface_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
     def load_chat_model(self):
-        if self.config["llm_params"]["llm_loader"] in ["gpt-3.5-turbo-1106", "gpt-4"]:
+        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.

code/modules/chat/langgraph/langgraph_rag.py

@@ -1,303 +0,0 @@
-# Adapted from https://github.com/langchain-ai/langgraph/blob/main/examples/rag/langgraph_crag.ipynb?ref=blog.langchain.dev
-
-from typing import List
-
-from typing_extensions import TypedDict
-from langgraph.graph import END, StateGraph, START
-from modules.chat.base import BaseRAG
-from langchain.memory import ChatMessageHistory
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.pydantic_v1 import BaseModel, Field
-from langchain_openai import ChatOpenAI
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.prompts import ChatPromptTemplate
-
-
-class GradeDocuments(BaseModel):
-    """Binary score for relevance check on retrieved documents."""
-
-    binary_score: str = Field(
-        description="Documents are relevant to the question, 'yes' or 'no'"
-    )
-
-
-class GraphState(TypedDict):
-    """
-    Represents the state of our graph.
-
-    Attributes:
-        question: question
-        generation: LLM generation
-        documents: list of documents
-    """
-
-    question: str
-    generation: str
-    documents: List[str]
-
-
-class Langgraph_RAG(BaseRAG):
-    def __init__(self, llm, memory, retriever, qa_prompt: str, rephrase_prompt: str):
-        """
-        Initialize the Langgraph_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.structured_llm_grader = llm.with_structured_output(GradeDocuments)
-        self.memory = self.add_history_from_list(memory)
-        self.retriever = retriever
-        self.qa_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. "
-            "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: {question}\n"
-            "AI Tutor:"
-        )
-        self.rephrase_prompt = rephrase_prompt
-        self.store = {}
-
-        ## Fix below ##
-
-        system = """You are a grader assessing relevance of a retrieved document to a user question. \n 
-            If the document contains keyword(s) or semantic meaning related to the question, grade it as relevant. \n
-            Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""
-        grade_prompt = ChatPromptTemplate.from_messages(
-            [
-                ("system", system),
-                (
-                    "human",
-                    "Retrieved document: \n\n {document} \n\n User question: {question}",
-                ),
-            ]
-        )
-
-        self.retrieval_grader = grade_prompt | self.structured_llm_grader
-
-        system = """You a question re-writer that converts an input question to a better version that is optimized \n 
-            for web search. Look at the input and try to reason about the underlying semantic intent / meaning."""
-        re_write_prompt = ChatPromptTemplate.from_messages(
-            [
-                ("system", system),
-                (
-                    "human",
-                    "Here is the initial question: \n\n {question} \n Formulate an improved question.",
-                ),
-            ]
-        )
-
-        self.question_rewriter = re_write_prompt | self.llm | StrOutputParser()
-
-        # Generate
-        self.qa_prompt_template = ChatPromptTemplate.from_template(self.qa_prompt)
-        self.rag_chain = self.qa_prompt_template | self.llm | StrOutputParser()
-
-        ###
-
-        # build the agentic graph
-        self.app = self.create_agentic_graph()
-
-    def retrieve(self, state):
-        """
-        Retrieve documents
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-            state (dict): New key added to state, documents, that contains retrieved documents
-        """
-        print("---RETRIEVE---")
-        question = state["question"]
-
-        # Retrieval
-        documents = self.retriever.get_relevant_documents(question)
-        return {"documents": documents, "question": question}
-
-    def generate(self, state):
-        """
-        Generate answer
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-            state (dict): New key added to state, generation, that contains LLM generation
-        """
-        print("---GENERATE---")
-        question = state["question"]
-        documents = state["documents"]
-
-        # RAG generation
-        generation = self.rag_chain.invoke({"context": documents, "question": question})
-        return {"documents": documents, "question": question, "generation": generation}
-
-    def transform_query(self, state):
-        """
-        Transform the query to produce a better question.
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-            state (dict): Updates question key with a re-phrased question
-        """
-
-        print("---TRANSFORM QUERY---")
-        question = state["question"]
-        documents = state["documents"]
-
-        # Re-write question
-        better_question = self.question_rewriter.invoke({"question": question})
-        return {"documents": documents, "question": better_question}
-
-    def grade_documents(self, state):
-        """
-        Determines whether the retrieved documents are relevant to the question.
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-            state (dict): Updates documents key with only filtered relevant documents
-        """
-
-        print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
-        question = state["question"]
-        documents = state["documents"]
-
-        # Score each doc
-        filtered_docs = []
-        web_search = "No"
-        for d in documents:
-            score = self.retrieval_grader.invoke(
-                {"question": question, "document": d.page_content}
-            )
-            grade = score.binary_score
-            if grade == "yes":
-                print("---GRADE: DOCUMENT RELEVANT---")
-                filtered_docs.append(d)
-            else:
-                print("---GRADE: DOCUMENT NOT RELEVANT---")
-                web_search = "Yes"
-                continue
-        return {
-            "documents": filtered_docs,
-            "question": question,
-            "web_search": web_search,
-        }
-
-    def decide_to_generate(self, state):
-        """
-        Determines whether to generate an answer, or re-generate a question.
-
-        Args:
-            state (dict): The current graph state
-
-        Returns:
-            str: Binary decision for next node to call
-        """
-
-        print("---ASSESS GRADED DOCUMENTS---")
-        state["question"]
-        web_search = state["web_search"]
-        state["documents"]
-
-        if web_search == "Yes":
-            # All documents have been filtered check_relevance
-            # We will re-generate a new query
-            print(
-                "---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---"
-            )
-            return "transform_query"
-        else:
-            # We have relevant documents, so generate answer
-            print("---DECISION: GENERATE---")
-            return "generate"
-
-    def create_agentic_graph(self):
-        """
-        Create an agentic graph to answer questions.
-
-        Returns:
-            dict: Agentic graph
-        """
-        self.workflow = StateGraph(GraphState)
-        self.workflow.add_node("retrieve", self.retrieve)
-        self.workflow.add_node(
-            "grade_documents", self.grade_documents
-        )  # grade documents
-        self.workflow.add_node("generate", self.generate)  # generatae
-        self.workflow.add_node(
-            "transform_query", self.transform_query
-        )  # transform_query
-
-        # build the graph
-        self.workflow.add_edge(START, "retrieve")
-        self.workflow.add_edge("retrieve", "grade_documents")
-        self.workflow.add_conditional_edges(
-            "grade_documents",
-            self.decide_to_generate,
-            {
-                "transform_query": "transform_query",
-                "generate": "generate",
-            },
-        )
-
-        self.workflow.add_edge("transform_query", "generate")
-        self.workflow.add_edge("generate", END)
-
-        # Compile
-        app = self.workflow.compile()
-        return app
-
-    def invoke(self, user_query, config):
-        """
-        Invoke the chain.
-
-        Args:
-            kwargs: The input variables.
-
-        Returns:
-            dict: The output variables.
-        """
-
-        inputs = {
-            "question": user_query["input"],
-        }
-
-        for output in self.app.stream(inputs):
-            for key, value in output.items():
-                # Node
-                print(f"Node {key} returned: {value}")
-            print("\n\n")
-
-        print(value["generation"])
-
-        # rename generation to answer
-        value["answer"] = value.pop("generation")
-        value["context"] = value.pop("documents")
-
-        return value
-
-    def add_history_from_list(self, history_list):
-        """
-        Add messages from a list to the chat history.
-
-        Args:
-            messages (list): The list of messages to add.
-        """
-        history = ChatMessageHistory()
-
-        for idx, message_pairs in enumerate(history_list):
-            history.add_user_message(message_pairs[0])
-            history.add_ai_message(message_pairs[1])
-
-        return history

code/modules/chat/llm_tutor.py

@@ -3,7 +3,6 @@ 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
-from modules.chat.langgraph.langgraph_rag import Langgraph_RAG
 
 
 class LLMTutor:
@@ -111,14 +110,6 @@ class LLMTutor:
                 qa_prompt=qa_prompt,
                 rephrase_prompt=rephrase_prompt,
             )
-        elif self.config["llm_params"]["llm_arch"] == "langgraph_agentic":
-            self.qa_chain = Langgraph_RAG(
-                llm=llm,
-                memory=memory,
-                retriever=retriever,
-                qa_prompt=qa_prompt,
-                rephrase_prompt=rephrase_prompt,
-            )
         else:
             raise ValueError(
                 f"Invalid LLM Architecture: {self.config['llm_params']['llm_arch']}"

code/modules/chat_processor/base.py

@@ -1,18 +0,0 @@
-# Template for chat processor classes
-
-
-class ChatProcessorBase:
-    def __init__(self):
-        pass
-
-    def process(self, message):
-        """
-        Processes and Logs the message
-        """
-        raise NotImplementedError("process method not implemented")
-
-    async def rag(self, user_query: dict, config: dict, chain):
-        """
-        Retrieves the response from the chain
-        """
-        raise NotImplementedError("rag method not implemented")

code/modules/chat_processor/chat_processor.py

@@ -1,55 +0,0 @@
-from modules.chat_processor.literal_ai import LiteralaiChatProcessor
-
-
-class ChatProcessor:
-    def __init__(self, config, user, tags=None):
-        self.config = config
-        self.chat_processor_type = self.config["chat_logging"]["platform"]
-        self.logging = self.config["chat_logging"]["log_chat"]
-        self.user = 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"])
-        return tags
-
-    def _init_processor(self):
-        if self.chat_processor_type == "literalai":
-            self.processor = LiteralaiChatProcessor(self.user, 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, stream):
-        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.config["llm_params"]["memory_window"],
-            }
-        }
-
-        # Process the user query using the appropriate method
-        if self.logging:
-            return await self.processor.rag(
-                user_query=user_query_dict, config=config, chain=chain
-            )
-        else:
-            if stream:
-                return chain.stream(user_query=user_query_dict, config=config)
-            return chain.invoke(user_query=user_query_dict, config=config)

code/modules/chat_processor/literal_ai.py

@@ -1,110 +1,7 @@
-from literalai import LiteralClient
-from literalai.api import LiteralAPI
-from literalai.filter import Filter as ThreadFilter
+from chainlit.data import ChainlitDataLayer
 
-import os
-from .base import ChatProcessorBase
 
-
-class LiteralaiChatProcessor(ChatProcessorBase):
-    def __init__(self, user=None, tags=None):
-        super().__init__()
-        self.user = user
-        self.tags = tags
-        self.literal_client = LiteralClient(api_key=os.getenv("LITERAL_API_KEY"))
-        self.literal_api = LiteralAPI(
-            api_key=os.getenv("LITERAL_API_KEY"), url=os.getenv("LITERAL_API_URL")
-        )
-        self.literal_client.reset_context()
-        self.user_info = self._fetch_userinfo()
-        self.user_thread = self._fetch_user_threads()
-        if len(self.user_thread["data"]) == 0:
-            self.thread = self._create_user_thread()
-        else:
-            self.thread = self._get_user_thread()
-        self.thread_id = self.thread["id"]
-
-        self.prev_conv = self._get_prev_k_conversations()
-
-    def _get_user_thread(self):
-        thread = self.literal_api.get_thread(id=self.user_thread["data"][0]["id"])
-        return thread.to_dict()
-
-    def _create_user_thread(self):
-        thread = self.literal_api.create_thread(
-            name=f"{self.user_info['identifier']}",
-            participant_id=self.user_info["metadata"]["id"],
-            environment="dev",
-        )
-
-        return thread.to_dict()
-
-    def _get_prev_k_conversations(self, k=3):
-
-        steps = self.thread["steps"]
-        conversation_pairs = []
-        count = 0
-        for i in range(len(steps) - 1, 0, -1):
-            if (
-                steps[i - 1]["type"] == "user_message"
-                and steps[i]["type"] == "assistant_message"
-            ):
-                user_message = steps[i - 1]["output"]["content"]
-                assistant_message = steps[i]["output"]["content"]
-                conversation_pairs.append((user_message, assistant_message))
-
-                count += 1
-                if count >= k:
-                    break
-
-        # Return the last k conversation pairs, reversed to maintain chronological order
-        return conversation_pairs[::-1]
-
-    def _fetch_user_threads(self):
-        filters = filters = [
-            {
-                "operator": "eq",
-                "field": "participantId",
-                "value": self.user_info["metadata"]["id"],
-            }
-        ]
-        user_threads = self.literal_api.get_threads(filters=filters)
-        return user_threads.to_dict()
-
-    def _fetch_userinfo(self):
-        user_info = self.literal_api.get_or_create_user(
-            identifier=self.user["user_id"]
-        ).to_dict()
-        # TODO: Have to do this more elegantly
-        # update metadata with unique id for now
-        # (literalai seems to not return the unique id as of now,
-        # so have to explicitly update it in the metadata)
-        user_info = self.literal_api.update_user(
-            id=user_info["id"],
-            metadata={
-                "id": user_info["id"],
-            },
-        ).to_dict()
-        return user_info
-
-    def process(self, user_message, assistant_message, source_dict):
-        with self.literal_client.thread(thread_id=self.thread_id) as thread:
-            self.literal_client.message(
-                content=user_message,
-                type="user_message",
-                name="User",
-            )
-            self.literal_client.message(
-                content=assistant_message,
-                type="assistant_message",
-                name="AI_Tutor",
-            )
-
-    async def rag(self, user_query: dict, config: dict, chain):
-        with self.literal_client.step(
-            type="retrieval", name="RAG", thread_id=self.thread_id, tags=self.tags
-        ) as step:
-            step.input = {"question": user_query["input"]}
-            res = chain.invoke(user_query, config)
-            step.output = res
-        return res
+# 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)

code/modules/config/config.yml

@@ -4,13 +4,12 @@ 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
   expand_urls: True # bool
-  db_option : 'FAISS' # str [FAISS, Chroma, RAGatouille, RAPTOR]
-  db_path : 'vectorstores' # str
+  db_option : 'RAGatouille' # str [FAISS, Chroma, RAGatouille, RAPTOR]
+  db_path : '../vectorstores' # str
   model : 'sentence-transformers/all-MiniLM-L6-v2' # str [sentence-transformers/all-MiniLM-L6-v2, text-embedding-ada-002']
   search_top_k : 3 # int
   score_threshold : 0.2 # float
@@ -30,7 +29,7 @@ llm_params:
   use_history: True # bool
   memory_window: 3 # int
   llm_style: 'Normal' # str [Normal, ELI5, Socratic]
-  llm_loader: 'gpt-3.5-turbo-1106' # str [local_llm, gpt-3.5-turbo-1106, gpt-4]
+  llm_loader: 'gpt-4o-mini' # str [local_llm, gpt-3.5-turbo-1106, gpt-4, gpt-4o-mini]
   openai_params:
     temperature: 0.7 # float
   local_llm_params:

code/modules/config/constants.py

@@ -7,7 +7,7 @@ load_dotenv()
 
 OPENAI_API_KEY = os.getenv("OPENAI_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")
 
 OAUTH_GOOGLE_CLIENT_ID = os.getenv("OAUTH_GOOGLE_CLIENT_ID")
@@ -18,3 +18,5 @@ opening_message = f"Hey, What Can I Help You With?\n\nYou can me ask me question
 # Model Paths
 
 LLAMA_PATH = "../storage/models/tinyllama-1.1b-chat-v1.0.Q5_K_M.gguf"
+
+RETRIEVER_HF_PATHS = {"RAGatouille": "XThomasBU/Colbert_Index"}

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,7 +136,13 @@ 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 {self.config['vectorstore']['db_option']} from Hugging Face: {end_time - start_time} seconds"
@@ -143,9 +150,9 @@ class VectorStoreManager:
         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 Download database {self.config['vectorstore']['db_option']} from Hugging Face: {end_time - start_time} seconds"
@@ -164,8 +171,14 @@ 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,
         )

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;
+} */