init commit for chainlit improvements

code/main.py

@@ -1,176 +1,244 @@
-from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
-from langchain_core.prompts import PromptTemplate
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain.chains import RetrievalQA
+import json
+import textwrap
+from typing import Any, Callable, Dict, List, Literal, Optional, no_type_check
 import chainlit as cl
-from langchain_community.chat_models import ChatOpenAI
-from langchain_community.embeddings import OpenAIEmbeddings
+from chainlit import run_sync
+from chainlit.config import config
 import yaml
-import logging
-from dotenv import load_dotenv
+import os
 
 from modules.chat.llm_tutor import LLMTutor
-from modules.config.constants import *
-from modules.chat.helpers import get_sources
 from modules.chat_processor.chat_processor import ChatProcessor
+from modules.config.constants import LLAMA_PATH
+from modules.chat.helpers import get_sources
 
-global logger
-# Initialize logger
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
-
-# Console Handler
-console_handler = logging.StreamHandler()
-console_handler.setLevel(logging.INFO)
-console_handler.setFormatter(formatter)
-logger.addHandler(console_handler)
-
-
-@cl.set_starters
-async def set_starters():
-    return [
-        cl.Starter(
-            label="recording on CNNs?",
-            message="Where can I find the recording for the lecture on Transfromers?",
-            icon="/public/adv-screen-recorder-svgrepo-com.svg",
-        ),
-        cl.Starter(
-            label="where's the slides?",
-            message="When are the lectures? I can't find the schedule.",
-            icon="/public/alarmy-svgrepo-com.svg",
-        ),
-        cl.Starter(
-            label="Due Date?",
-            message="When is the final project due?",
-            icon="/public/calendar-samsung-17-svgrepo-com.svg",
-        ),
-        cl.Starter(
-            label="Explain backprop.",
-            message="I didnt understand the math behind backprop, could you explain it?",
-            icon="/public/acastusphoton-svgrepo-com.svg",
-        ),
-    ]
-
-
-# Adding option to select the chat profile
-@cl.set_chat_profiles
-async def chat_profile():
-    return [
-        # cl.ChatProfile(
-        #     name="Mistral",
-        #     markdown_description="Use the local LLM: **Mistral**.",
-        # ),
-        cl.ChatProfile(
-            name="gpt-3.5-turbo-1106",
-            markdown_description="Use OpenAI API for **gpt-3.5-turbo-1106**.",
-        ),
-        cl.ChatProfile(
-            name="gpt-4",
-            markdown_description="Use OpenAI API for **gpt-4**.",
-        ),
-        cl.ChatProfile(
-            name="Llama",
-            markdown_description="Use the local LLM: **Tiny Llama**.",
-        ),
-    ]
-
-
-@cl.author_rename
-def rename(orig_author: str):
-    rename_dict = {"Chatbot": "AI Tutor"}
-    return rename_dict.get(orig_author, orig_author)
-
-
-# chainlit code
-@cl.on_chat_start
-async def start():
-    with open("modules/config/config.yml", "r") as f:
-        config = yaml.safe_load(f)
-
-    # Ensure log directory exists
-    log_directory = config["log_dir"]
-    if not os.path.exists(log_directory):
-        os.makedirs(log_directory)
-
-    # File Handler
-    log_file_path = (
-        f"{log_directory}/tutor.log"  # Change this to your desired log file path
-    )
-    file_handler = logging.FileHandler(log_file_path, mode="w")
-    file_handler.setLevel(logging.INFO)
-    file_handler.setFormatter(formatter)
-    logger.addHandler(file_handler)
-
-    logger.info("Config file loaded")
-    logger.info(f"Config: {config}")
-    logger.info("Creating llm_tutor instance")
-
-    chat_profile = cl.user_session.get("chat_profile")
-    if chat_profile is not None:
-        if chat_profile.lower() in ["gpt-3.5-turbo-1106", "gpt-4"]:
-            config["llm_params"]["llm_loader"] = "openai"
-            config["llm_params"]["openai_params"]["model"] = chat_profile.lower()
-        elif chat_profile.lower() == "llama":
-            config["llm_params"]["llm_loader"] = "local_llm"
-            config["llm_params"]["local_llm_params"]["model"] = LLAMA_PATH
-            config["llm_params"]["local_llm_params"]["model_type"] = "llama"
-        elif chat_profile.lower() == "mistral":
-            config["llm_params"]["llm_loader"] = "local_llm"
-            config["llm_params"]["local_llm_params"]["model"] = MISTRAL_PATH
-            config["llm_params"]["local_llm_params"]["model_type"] = "mistral"
-
-        else:
-            pass
-
-    llm_tutor = LLMTutor(config, logger=logger)
-
-    chain = llm_tutor.qa_bot()
-    # msg = cl.Message(content=f"Starting the bot {chat_profile}...")
-    # await msg.send()
-    # msg.content = opening_message
-    # await msg.update()
-
-    tags = [chat_profile, config["vectorstore"]["db_option"]]
-    chat_processor = ChatProcessor(config, tags=tags)
-    cl.user_session.set("chain", chain)
-    cl.user_session.set("counter", 0)
-    cl.user_session.set("chat_processor", chat_processor)
-
-
-@cl.on_chat_end
-async def on_chat_end():
-    await cl.Message(content="Sorry, I have to go now. Goodbye!").send()
-
-
-@cl.on_message
-async def main(message):
-    global logger
-    user = cl.user_session.get("user")
-    chain = cl.user_session.get("chain")
-
-    counter = cl.user_session.get("counter")
-    counter += 1
-    cl.user_session.set("counter", counter)
-
-    # if counter >= 3:  # Ensure the counter condition is checked
-    #     await cl.Message(content="Your credits are up!").send()
-    #     await on_chat_end()  # Call the on_chat_end function to handle the end of the chat
-    #     return  # Exit the function to stop further processing
-    # else:
-
-    cb = cl.AsyncLangchainCallbackHandler()  # TODO: fix streaming here
-    cb.answer_reached = True
-
-    processor = cl.user_session.get("chat_processor")
-    res = await processor.rag(message.content, chain, cb)
-    try:
-        answer = res["answer"]
-    except:
-        answer = res["result"]
-
-    answer_with_sources, source_elements, sources_dict = get_sources(res, answer)
-    processor._process(message.content, answer, sources_dict)
-
-    await cl.Message(content=answer_with_sources, elements=source_elements).send()
+from chainlit.input_widget import Select, Switch, Slider
+
+USER_TIMEOUT = 60_000
+SYSTEM = "System 🖥️"
+LLM = "LLM 🧠"
+AGENT = "Agent <>"
+YOU = "You 😃"
+ERROR = "Error 🚫"
+
+
+class Chatbot:
+    def __init__(self):
+        self.llm_tutor = None
+        self.chain = None
+        self.chat_processor = None
+        self.config = self._load_config()
+
+    def _load_config(self):
+        with open("modules/config/config.yml", "r") as f:
+            config = yaml.safe_load(f)
+        return config
+
+    async def ask_helper(func, **kwargs):
+        res = await func(**kwargs).send()
+        while not res:
+            res = await func(**kwargs).send()
+        return res
+
+    @no_type_check
+    async def setup_llm(self) -> None:
+        """From the session `llm_settings`, create new LLMConfig and LLM objects,
+        save them in session state."""
+
+        llm_settings = cl.user_session.get("llm_settings", {})
+        chat_profile = llm_settings.get("chat_model")
+        retriever_method = llm_settings.get("retriever_method")
+        memory_window = llm_settings.get("memory_window")
+
+        self._configure_llm(chat_profile)
+
+        chain = cl.user_session.get("chain")
+        memory = chain.memory
+        self.config["vectorstore"][
+            "db_option"
+        ] = retriever_method  # update the retriever method in the config
+        memory.k = memory_window  # set the memory window
+
+        self.llm_tutor = LLMTutor(self.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)
+
+        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]) -> None:
+        """Update LLMConfig and LLM from settings, and save in session state."""
+        cl.user_session.set("llm_settings", new_settings)
+        await self.inform_llm_settings()
+        await self.setup_llm()
+
+    async def make_llm_settings_widgets(self, config=None):
+        config = config or self.config
+        await cl.ChatSettings(
+            [
+                cl.input_widget.Select(
+                    id="chat_model",
+                    label="Model Name (Default GPT-3)",
+                    values=["llama", "gpt-3.5-turbo-1106", "gpt-4"],
+                    initial_index=0,
+                ),
+                cl.input_widget.Select(
+                    id="retriever_method",
+                    label="Retriever (Default FAISS)",
+                    values=["FAISS", "Chroma", "RAGatouille", "RAPTOR"],
+                    initial_index=0,
+                ),
+                cl.input_widget.Slider(
+                    id="memory_window",
+                    label="Memory Window (Default 3)",
+                    initial=3,
+                    min=0,
+                    max=10,
+                    step=1,
+                ),
+                cl.input_widget.Switch(
+                    id="view_sources", label="View Sources", initial=False
+                ),
+            ]
+        ).send()  # type: ignore
+
+    @no_type_check
+    async def inform_llm_settings(self) -> None:
+        llm_settings: Dict[str, Any] = cl.user_session.get("llm_settings", {})
+        llm_tutor = cl.user_session.get("llm_tutor")
+        settings_dict = dict(
+            model=llm_settings.get("chat_model"),
+            retriever=llm_settings.get("retriever_method"),
+            memory_window=llm_settings.get("memory_window"),
+            num_docs_in_db=len(llm_tutor.vector_db),
+            view_sources=llm_settings.get("view_sources"),
+        )
+        await cl.Message(
+            author=SYSTEM,
+            content="LLM settings have been updated. You can continue with your Query!",
+            elements=[
+                cl.Text(
+                    name="settings",
+                    display="side",
+                    content=json.dumps(settings_dict, indent=4),
+                    language="json",
+                )
+            ],
+        ).send()
+
+    async def set_starters(self):
+        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",
+            ),
+        ]
+
+    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="Llama",
+                markdown_description="Use the local LLM: **Tiny Llama**.",
+            ),
+        ]
+
+    def rename(self, orig_author: str):
+        rename_dict = {"Chatbot": "AI Tutor"}
+        return rename_dict.get(orig_author, orig_author)
+
+    async def start(self):
+        await self.make_llm_settings_widgets(self.config)
+
+        chat_profile = cl.user_session.get("chat_profile")
+        if chat_profile:
+            self._configure_llm(chat_profile)
+
+        self.llm_tutor = LLMTutor(self.config)
+        self.chain = self.llm_tutor.qa_bot()
+        tags = [chat_profile, self.config["vectorstore"]["db_option"]]
+        self.chat_processor = ChatProcessor(self.config, 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("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")
+
+        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)
+        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")
+        )
+        processor._process(message.content, answer, sources_dict)
+
+        await cl.Message(content=answer_with_sources, elements=source_elements).send()
+
+    def _configure_llm(self, chat_profile):
+        chat_profile = chat_profile.lower()
+        if chat_profile in ["gpt-3.5-turbo-1106", "gpt-4"]:
+            self.config["llm_params"]["llm_loader"] = "openai"
+            self.config["llm_params"]["openai_params"]["model"] = chat_profile
+        elif chat_profile == "llama":
+            self.config["llm_params"]["llm_loader"] = "local_llm"
+            self.config["llm_params"]["local_llm_params"]["model"] = LLAMA_PATH
+            self.config["llm_params"]["local_llm_params"]["model_type"] = "llama"
+        elif chat_profile == "mistral":
+            self.config["llm_params"]["llm_loader"] = "local_llm"
+            self.config["llm_params"]["local_llm_params"]["model"] = MISTRAL_PATH
+            self.config["llm_params"]["local_llm_params"]["model_type"] = "mistral"
+
+
+chatbot = Chatbot()
+
+# Register functions to Chainlit events
+cl.set_starters(chatbot.set_starters)
+cl.set_chat_profiles(chatbot.chat_profile)
+cl.author_rename(chatbot.rename)
+cl.on_chat_start(chatbot.start)
+cl.on_chat_end(chatbot.on_chat_end)
+cl.on_message(chatbot.main)
+cl.on_settings_update(chatbot.update_llm)

code/modules/chat/helpers.py

@@ -3,7 +3,7 @@ import chainlit as cl
 from langchain_core.prompts import PromptTemplate
 
 
-def get_sources(res, answer):
+def get_sources(res, answer, view_sources=False):
     source_elements = []
     source_dict = {}  # Dictionary to store URL elements
 
@@ -40,40 +40,42 @@ def get_sources(res, answer):
     full_answer = "**Answer:**\n"
     full_answer += answer
 
-    # Then, display the sources
-    full_answer += "\n\n**Sources:**\n"
-    for idx, (url_name, source_data) in enumerate(source_dict.items()):
-        full_answer += f"\nSource {idx + 1} (Score: {source_data['score']}): {source_data['url']}\n"
+    if view_sources:
 
-        name = f"Source {idx + 1} Text\n"
-        full_answer += name
-        source_elements.append(
-            cl.Text(name=name, content=source_data["text"], display="side")
-        )
+        # 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"
 
-        # 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"))
+            source_elements.append(
+                cl.Text(name=name, content=source_data["text"], 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",
+            # 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
 

code/modules/chat/llm_tutor.py

@@ -157,18 +157,18 @@ class LLMTutor:
         return prompt
 
     # Retrieval QA Chain
-    def retrieval_qa_chain(self, llm, prompt, db):
+    def retrieval_qa_chain(self, llm, prompt, db, memory=None):
 
         retriever = Retriever(self.config)._return_retriever(db)
 
         if self.config["llm_params"]["use_history"]:
-            memory = ConversationBufferWindowMemory(
-                k=self.config["llm_params"]["memory_window"],
-                memory_key="chat_history",
-                return_messages=True,
-                output_key="answer",
-                max_token_limit=128,
-            )
+            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(
                 llm=llm,
                 chain_type="stuff",
@@ -195,11 +195,16 @@ class LLMTutor:
         return llm
 
     # QA Model Function
-    def qa_bot(self):
+    def qa_bot(self, memory=None):
         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
+            self.llm, qa_prompt, db, memory
         )  # TODO:  PROMPT is overwritten in CustomConversationalRetrievalChain
 
         return qa

code/modules/vectorstore/base.py

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

code/modules/vectorstore/chroma.py

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

code/modules/vectorstore/colbert.py

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

code/modules/vectorstore/faiss.py

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

code/modules/vectorstore/raptor.py

@@ -16,6 +16,13 @@ from modules.vectorstore.base import VectorStoreBase
 RANDOM_SEED = 42
 
 
+class FAISS(FAISS):
+    """To add length property to FAISS class"""
+
+    def __len__(self):
+        return self.index.ntotal
+
+
 class RAPTORVectoreStore(VectorStoreBase):
     def __init__(self, config, documents=[], text_splitter=None, embedding_model=None):
         self.documents = documents

code/modules/vectorstore/store_manager.py

@@ -138,7 +138,7 @@ class VectorStoreManager:
         self.loaded_vector_db = self.vector_db._load_database(self.embedding_model)
         end_time = time.time()  # End time for loading database
         self.logger.info(
-            f"Time taken to load database: {end_time - start_time} seconds"
+            f"Time taken to load database {self.config['vectorstore']['db_option']} from Hugging Face: {end_time - start_time} seconds"
         )
         self.logger.info("Loaded database")
         return self.loaded_vector_db
@@ -148,8 +148,12 @@ class VectorStoreManager:
         self.vector_db._load_from_HF()
         end_time = time.time()
         self.logger.info(
-            f"Time taken to load database from Hugging Face: {end_time - start_time} seconds"
+            f"Time taken to Download database {self.config['vectorstore']['db_option']} from Hugging Face: {end_time - start_time} seconds"
         )
+        self.logger.info("Downloaded database")
+
+    def __len__(self):
+        return len(self.vector_db)
 
 
 if __name__ == "__main__":

code/modules/vectorstore/vectorstore.py

@@ -86,3 +86,6 @@ class VectorStore:
 
     def _get_vectorstore(self):
         return self.vectorstore
+
+    def __len__(self):
+        return self.vectorstore.__len__()