Update app.py

app.py

@@ -1,75 +1,133 @@
-import chainlit as cl 
-from langchain_together import ChatTogether, TogetherEmbeddings
-from langchain_core.runnables import RunnableSequence, RunnablePassthrough
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_community.document_loaders import YoutubeLoader
-from typing import List
-import langchain_core
-from langchain_community.vectorstores import FAISS
-from langchain.retrievers.ensemble import EnsembleRetriever
-from langchain_community.retrievers import BM25Retriever
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+import chainlit as cl # handles the chat interface
+from langchain_together import ChatTogether, TogetherEmbeddings # for the LLM and Embeddings
+from langchain_core.runnables import RunnableSequence, RunnablePassthrough # for chain execution
+from langchain_core.prompts import ChatPromptTemplate # for writing the prompt template
+from langchain_community.document_loaders import YoutubeLoader # for loading the youtube video
+from typing import List # for type hinting
+import langchain_core # for type hinting
+from langchain_community.vectorstores import FAISS # for the vector store
+from langchain_community.retrievers import BM25Retriever # for the BM25 retriever
+from langchain.retrievers.ensemble import EnsembleRetriever # for the ensemble retriever
+from langchain_text_splitters import RecursiveCharacterTextSplitter # for the text splitter
 
-def create_youtube_transcription(youtube_url: str):
+def create_youtube_transcription(youtube_url: str) -> List[langchain_core.documents.Document]:
+    """
+    Create a youtube transcription from a youtube url
+    More Info: https://python.langchain.com/docs/integrations/document_loaders/youtube_transcript/
+    Accepts:
+        youtube_url: str - The url of the youtube video
+    Returns:
+        List[langchain_core.documents.Document]: A list of documents containing the youtube transcription
+    """
     loader = YoutubeLoader.from_youtube_url(
-    youtube_url, add_video_info=False
-    )
-    youtube_docs = loader.load()
+        youtube_url, add_video_info=True
+    ) # we can also pass an array of youtube urls to load multiple videos at once!
+    youtube_docs = loader.load() # this loads the transcript
     return youtube_docs
 
 def create_text_splitter(docs: List[langchain_core.documents.Document]):
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    docs = text_splitter.split_documents(docs)
+    """
+    Create a text splitter from a list of documents
+    More Info: ument_transformers/recursive_text_splitter/
+    Accepts:
+        docs: List[langchain_core.documents.Document] - A list of documents to split
+    Returns:
+        List[langchain_core.documents.Document]: A list of documents containing the text split
+    """
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) # without an overlap, context might get cut off
+    docs = text_splitter.split_documents(docs) # split the documents into chunks
     return docs
 
-def create_vector_store(docs: List[langchain_core.documents.Document]):
-    embedding = cl.user_session.get("embedding")
-    vector_db = FAISS.from_documents(docs, embedding)
+def create_faiss_vector_store(docs: List[langchain_core.documents.Document]) -> FAISS:
+    """
+    Create a FAISS vector store or vector database from a list of documents
+    More Info: https://python.langchain.com/docs/integrations/vectorstores/faiss/
+    Accepts:
+        docs: List[langchain_core.documents.Document] - A list of documents to store
+    Returns:
+        FAISS: A vector store containing the documents
+    """
+    embedding = cl.user_session.get("embedding") # we can get the embedding model from the user session or pass as a parameter too!
+    vector_db = FAISS.from_documents(docs, embedding) # create the vector store
+    vector_db.k = 5 # we set k to 5, so we get 5 documents back
     return vector_db
 
-def create_bm25_vector_store(docs: List[langchain_core.documents.Document]):
-    bm25 = BM25Retriever.from_documents(docs)
+def create_bm25_retreiver(docs: List[langchain_core.documents.Document]) -> BM25Retriever:
+    """
+    Create a BM25 retriever from a list of documents
+    More Info: https://python.langchain.com/docs/integrations/retrievers/bm25/
+    Accepts:
+        docs: List[langchain_core.documents.Document] - A list of documents to store
+    Returns:
+        BM25Retriever: A BM25 retriever containing the documents
+    """
+    bm25 = BM25Retriever.from_documents(docs) # we don't need embeddings for BM25, as it uses keyword matching!
+    bm25.k = 5 # we set k to 5, so we get 5 documents back
     return bm25
 
-def create_ensemble_retriever(vector_db:FAISS, bm25:BM25Retriever):
-    ensemble_retreiver = EnsembleRetriever(retrievers=[vector_db.as_retriever(), bm25], weights=[.3, .7])
+def create_ensemble_retriever(vector_db:FAISS, bm25:BM25Retriever) -> EnsembleRetriever:
+    """
+    Create an ensemble retriever from a vector db and a BM25 retriever
+    More Info: https://python.langchain.com/docs/how_to/ensemble_retriever/
+    Accepts:
+        vector_db: FAISS - A vector db
+        bm25: BM25Retriever - A BM25 retriever
+    Returns:
+        EnsembleRetriever: An ensemble retriever containing the vector db and the BM25 retriever
+    """
+    ensemble_retreiver = EnsembleRetriever(retrievers=[vector_db.as_retriever(), bm25], weights=[.3, .7]) # 30% semantic, 70% keyword retrieval
     return ensemble_retreiver
 
 @cl.on_chat_start
 async def start():
-    await cl.Message(content="my name is josh!").send()
-    try: 
-        llm = ChatTogether(model="meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo")
-        await cl.Message(content=f"model is successfully loaded").send()
-        cl.user_session.set("llm", llm)
-        embedding = TogetherEmbeddings(model="togethercomputer/m2-bert-80M-8k-retrieval")
-        cl.user_session.set("embedding", embedding)
+    """
+    More info: https://docs.chainlit.io/api-reference/lifecycle-hooks/on-chat-start
+    This function is called when the chat starts. Under the hood it handles all the complicated stuff for loading the UI. 
+    We explicitly load the model, embeddings, and retrievers.
+    Asks the user to provide the YouTube video link and loads the transcription.
+    With the transcription, it creates a vector store and a BM25 vector store. That is used to create an ensemble retriever combining the two.
+    """
+    await cl.Message(content="Hello! I am your AI assistant. I can help you with your questions about the video you provide.").send()
+    try: # a try catch block prevents the app from crashing if have an error
+        llm = ChatTogether(model="meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo") # initialize the LLM model
+        await cl.Message(content=f"model is successfully loaded").send() # we can send messages to be displayed with cl.Message().send()
+        cl.user_session.set("llm", llm) # we can store variables in a special memory called the user session, so we can use them in our on message function and more
+        embedding = TogetherEmbeddings(model="togethercomputer/m2-bert-80M-8k-retrieval") # initialize the embedding model
+        cl.user_session.set("embedding", embedding) # store the embedding model in the user session
         await cl.Message(content="embedding model loaded").send()
-        youtube_link = await cl.AskUserMessage("Please provide the YouTube video link").send()
-        youtube_docs = create_youtube_transcription(youtube_link['output'])
-        split_docs = create_text_splitter(youtube_docs)
-        vector_db = create_vector_store(split_docs)
-        bm25 = create_bm25_vector_store(split_docs)
-        ensemble_retriever = create_ensemble_retriever(vector_db, bm25)
-        cl.user_session.set("ensemble_retriever", ensemble_retriever)
-        transcription = youtube_docs[0].page_content
-        await cl.Message(content=f"youtube docs: {transcription}").send()
-        cl.user_session.set("transcription", transcription)
+        youtube_link = await cl.AskUserMessage("Please provide the YouTube video link").send() # We can ask the user for input using cl.AskUserMessage().send() which does not affect cl.on_message()
+        # more on ask user message: https://docs.chainlit.io/api-reference/ask/ask-for-input
+        await cl.Message(content=f"youtube link: {youtube_link}").send() # display and double check to make sure the link is correct
+        youtube_docs = create_youtube_transcription(youtube_link['content']) # create the youtube transcription
+        split_docs = create_text_splitter(youtube_docs) # split the documents into chunks
+        vector_db = create_faiss_vector_store(split_docs) # create the vector db
+        bm25 = create_bm25_retreiver(split_docs) # create the BM25 retreiver
+        ensemble_retriever = create_ensemble_retriever(vector_db, bm25) # create the ensemble retriever
+        cl.user_session.set("ensemble_retriever", ensemble_retriever) # store the ensemble retriever in the user session for our on message function
+        transcription = youtube_docs[0].page_content # get the transcription of the first document
+        await cl.Message(content=f"youtube docs: {transcription}").send() # display the transcription of the first document to show that we have the correct data
     except Exception as e:
-        await cl.Message(content=f"failed to load model: {e}").send()
+        await cl.Message(content=f"failed to load model: {e}").send() # display the error if we failed to load the model
 
 
 @cl.on_message
 async def message(message: cl.Message):
+    """
+    More info: https://docs.chainlit.io/api-reference/lifecycle-hooks/on-message
+    This function is called when the user sends a message. It uses the ensemble retriever to find the most relevant documents and feeds them into the LLM.
+    We can then display the answer and the relevant documents to the user.
+    """
     prompt_template = ChatPromptTemplate.from_template(template="""
-    You are a helpful assistant that can answer questions about the following video. Here is the appropriate chunks of context: {context}.
-    Answer the question: {question} but do not use any information outside of the video. Site the source or information you used to answer the question
-    """)
-    llm = cl.user_session.get("llm")
-    vector_db = cl.user_session.get("vector_db")
-    transcription = cl.user_session.get("transcription")
-    ensemble_retriever = cl.user_session.get("ensemble_retriever")
-    rag_chain = RunnableSequence({"context": ensemble_retriever, "question": RunnablePassthrough()}, prompt_template | llm)
-    response = rag_chain.invoke(message.content)
-    await cl.Message(content=response.content).send()
-
+        You are a helpful assistant that can answer questions about the following video. Here is the appropriate chunks of context: {context}.
+        Answer the question: {question} but do not use any information outside of the video. Site the source or information you used to answer the question
+    """) # we create a prompt template that we will use to format our prompt
+    llm = cl.user_session.get("llm") # we get the LLM model we initialized in the start function
+    ensemble_retriever = cl.user_session.get("ensemble_retriever") # we get the ensemble retriever we initialized in the start function
+    relevant_docs = ensemble_retriever.invoke(message.content) # we use the ensemble retriever to find the most relevant documents
+    cl.Message(content=f"Displaying Relevant Docs").send() # we display the relevant documents to the user
+    for doc in relevant_docs: # loop through the relevant documents and display each one!
+        await cl.Message(content=doc.page_content).send() 
+    rag_chain = RunnableSequence({"context": ensemble_retriever, "question": RunnablePassthrough()} | prompt_template | llm)
+    response = rag_chain.invoke(message.content) # we invoke the rag chain with the user's message
+    await cl.Message(content="Done Displaying Relevant Docs").send()
+    await cl.Message(content=f"LLM Response: {response.content}").send() # we display the response to the user