Update app.py

app.py

@@ -10,6 +10,66 @@ from langchain_community.retrievers import BM25Retriever # for the BM25 retrieve
 from langchain.retrievers.ensemble import EnsembleRetriever # for the ensemble retriever
 from langchain_text_splitters import RecursiveCharacterTextSplitter # for the text splitter
 
+
+
+######## Chainlit ########
+@cl.on_chat_start
+async def start():
+    """
+    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() # 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['content']}").send() # display and double check to make sure the link is correct
+        youtube_docs = await create_youtube_transcription(youtube_link['content']) # create the youtube transcription
+        transcription = youtube_docs # 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
+        split_docs = await create_text_splitter(youtube_docs) # split the documents into chunks
+        vector_db = await create_faiss_vector_store(split_docs) # create the vector db
+        bm25 = await create_bm25_retreiver(split_docs) # create the BM25 retreiver
+        ensemble_retriever = await 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
+    except Exception as e:
+        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
+    """) # 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() 
+    await cl.Message(content="Done Displaying Relevant Docs").send()
+    # question -> retrieve relevant docs -> format the question and context and add it to the prompt template -> pass to LLM 
+    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=f"LLM Response: {response.content}").send() # we display the response to the user
+
+######## Youtube ########
+
 async def create_youtube_transcription(youtube_url: str) -> List[langchain_core.documents.Document]:
     """
     Create a youtube transcription from a youtube url
@@ -28,7 +88,10 @@ async def create_youtube_transcription(youtube_url: str) -> List[langchain_core.
     except Exception as e:
         await cl.Message(content=f"failed to load youtube video: {e} Please refresh the page").send() # display the error if we failed to load the youtube video
 
-async def create_text_splitter(docs: List[langchain_core.documents.Document]):
+
+######## RAG ########
+
+async def create_text_splitter(docs: List[langchain_core.documents.Document]) -> List[langchain_core.documents.Document]:
     """
     Create a text splitter from a list of documents
     More Info: ument_transformers/recursive_text_splitter/
@@ -88,58 +151,4 @@ async def create_ensemble_retriever(vector_db:FAISS, bm25:BM25Retriever) -> Ense
         ensemble_retreiver = EnsembleRetriever(retrievers=[vector_db.as_retriever(), bm25], weights=[.3, .7]) # 30% semantic, 70% keyword retrieval
         return ensemble_retreiver
     except Exception as e:
-        await cl.Message(content=f"failed to create ensemble retriever: {e}").send() # display the error if we failed to create the ensemble retriever
-
-@cl.on_chat_start
-async def start():
-    """
-    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() # 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['content']}").send() # display and double check to make sure the link is correct
-        youtube_docs = await create_youtube_transcription(youtube_link['content']) # create the youtube transcription
-        transcription = youtube_docs # 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
-        split_docs = await create_text_splitter(youtube_docs) # split the documents into chunks
-        vector_db = await create_faiss_vector_store(split_docs) # create the vector db
-        bm25 = await create_bm25_retreiver(split_docs) # create the BM25 retreiver
-        ensemble_retriever = await 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
-    except Exception as e:
-        await cl.Message(content=f"error that happened: {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
-    """) # 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() 
-    await cl.Message(content="Done Displaying Relevant Docs").send()
-    # question -> retrieve relevant docs -> format the question and context and add it to the prompt template -> pass to LLM 
-    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=f"LLM Response: {response.content}").send() # we display the response to the user
+        await cl.Message(content=f"failed to create ensemble retriever: {e}").send() # display the error if we failed to create the ensemble retriever