updated ui added sidebar

app.py

@@ -15,18 +15,34 @@ from langchain.output_parsers import ResponseSchema, StructuredOutputParser
 from langchain_core.runnables.history import RunnableWithMessageHistory
 from langchain_core.chat_history import BaseChatMessageHistory
 from langchain.chains import RetrievalQA
-# import nest_asyncio
-# nest_asyncio.apply()
-
-st.title('ASMR Query Bot 🔔')
+import nest_asyncio
+import pymongo
+import logging
+from langchain.docstore.document import Document
+import redis
+import threading
 
 # config
+nest_asyncio.apply()
+logging.basicConfig(level=logging.INFO)
 database = "AlertSimAndRemediation"
 collection = "alert_embed"
 index_name = "alert_index"
+stream_name = "alerts"
+redis_port = 16652
+
+# Streamlit Application
+st.set_page_config(
+    page_title="ASMR Query Bot 🔔",
+    page_icon="🔔",
+    layout="wide",
+    initial_sidebar_state="auto",
+    menu_items={
+        'About': "https://github.com/ankush-003/alerts-simulation-and-remediation"
+    }
+)
 
-# llm
-chat = ChatGroq(temperature=0, model_name="mixtral-8x7b-32768")
+st.title('ASMR Query Bot 🔔')
 
 # embedding model
 embedding_args = {
@@ -36,9 +52,6 @@ embedding_args = {
 }
 embedding_model = HuggingFaceEmbeddings(**embedding_args)
 
-# chat history
-# chat_history = ChatMessageHistory()
-
 # vector search
 vector_search = MongoDBAtlasVectorSearch.from_connection_string(
     os.environ["MONGO_URI"],
@@ -47,11 +60,6 @@ vector_search = MongoDBAtlasVectorSearch.from_connection_string(
     index_name=index_name,
 )
 
-qa_retriever = vector_search.as_retriever(
-    search_type="similarity",
-    search_kwargs={"k": 5},
-)
-
 # contextualising prev chats
 contextualize_q_system_prompt = """Given a chat history and the latest user question \
 which might reference context in the chat history, formulate a standalone question \
@@ -64,9 +72,6 @@ contextualize_q_prompt = ChatPromptTemplate.from_messages(
         ("human", "{input}"),
     ]
 )
-history_aware_retriever = create_history_aware_retriever(
-    chat, qa_retriever, contextualize_q_prompt
-)
 
 # prompt
 system_prompt = """
@@ -79,8 +84,6 @@ Your responses should be clear, concise, and tailored to the specific alert deta
 </context>
 """
 
-chat_history = []
-
 qa_prompt = ChatPromptTemplate.from_messages(
     [
         ("system", system_prompt),
@@ -88,58 +91,6 @@ qa_prompt = ChatPromptTemplate.from_messages(
         ("human", "{input}"),
     ]
 )
-question_answer_chain = create_stuff_documents_chain(chat, qa_prompt)
-
-# output parser
-response_schemas = [
-    ResponseSchema(name="answer", description="answer to the user's question"),
-    ResponseSchema(
-        name="source",
-        description="source used to answer the user's question, should be a website.",
-    )
-]
-output_parser = StructuredOutputParser.from_response_schemas(response_schemas)
-
-
-rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
-
-# managing message history
-# store = {}
-
-# def get_session_history(session_id: str) -> BaseChatMessageHistory:
-#     if session_id not in store:
-#         store[session_id] = ChatMessageHistory()
-#     return store[session_id]
-
-
-
-# conversational_rag_chain = RunnableWithMessageHistory(
-#     rag_chain,
-#     get_session_history,
-#     input_messages_key="input",
-#     history_messages_key="chat_history",
-#     output_messages_key="answer",
-# )
-
-# schema
-# print(conversational_rag_chain.input_schema.schema())
-# print(conversational_rag_chain.output_schema.schema())
-
-
-# Retrieves documents
-# retriever_chain = create_history_aware_retriever(chat, qa_retriever, prompt)
-
-# retriever_chain.invoke({
-#     "chat_history": chat_history,
-#     "input": "Tell me about the latest alert"
-# })
-
-# conversational_rag_chain.invoke(
-#     {"input": "What is the remedy to the latest alert"},
-#     config={
-#         "configurable": {"session_id": "abc123"}
-#     },  # constructs a key "abc123" in `store`.
-# )
 
 if "chat_messages" not in st.session_state:
     st.session_state.chat_messages = []
@@ -148,27 +99,75 @@ if "chat_messages" not in st.session_state:
 history = StreamlitChatMessageHistory(key="chat_messages")
 
 # Initialize chat history
-
-
 if len(history.messages) == 0:
-    history.add_ai_message("How can I help you?")
-
-conversational_rag_chain = RunnableWithMessageHistory(
-    rag_chain,
-    lambda session_id: history,
-    input_messages_key="input",
-    history_messages_key="chat_history",
-    output_messages_key="answer",
-)
+    history.add_ai_message("Hey I am ASMR Query Bot, how can i help you ?")
+
+with st.sidebar:
+    st.title('Settings ⚙️')
+    st.subheader('Models and parameters')
+    selected_model = st.sidebar.selectbox('Choose a model', ['Llama3-8B', 'Llama3-70B', 'Mixtral-8x7B'], key='selected_model')
+    if selected_model == 'Mixtral-8x7B':
+        model_name="mixtral-8x7b-32768"
+    elif selected_model == 'Llama3-70B':
+        model_name='Llama3-70b-8192'
+    elif selected_model == 'Llama3-8B':
+        model_name='Llama3-8b-8192'
+    temp = st.sidebar.slider('temperature', min_value=0.01, max_value=1.0, value=0.0, step=0.01)
+    k = st.sidebar.slider('number of docs retrieved', min_value=1, max_value=20, value=2, step=1)
+
+def get_response(query, config):
+  chat = ChatGroq(temperature=temp, model_name=model_name)
+  qa_retriever = vector_search.as_retriever(
+      search_type="similarity",
+      search_kwargs={"k": k},
+  )
+  history_aware_retriever = create_history_aware_retriever(
+      chat, qa_retriever, contextualize_q_prompt
+  )
+  question_answer_chain = create_stuff_documents_chain(chat, qa_prompt)
+  rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
+  conversational_rag_chain = RunnableWithMessageHistory(
+      rag_chain,
+      lambda session_id: history,
+      input_messages_key="input",
+      history_messages_key="chat_history",
+      output_messages_key="answer",
+  )
+  return conversational_rag_chain.invoke({"input": prompt}, config=config)
+
+def clear_chat_history():
+    st.session_state.chat_messages = []
+    history.add_ai_message("Hey I am ASMR Query Bot, how can i help you ?")
+
+st.sidebar.button('Clear Chat History', on_click=clear_chat_history)  
 
 for msg in history.messages:
     st.chat_message(msg.type).write(msg.content)
 
+# preprocessing context
+def format_docs_with_metadata(docs):
+    formatted_docs = []
+    for i, doc in enumerate(docs, start=1):
+        metadata_str = "\n".join([f"**{key}**: `{value}`\n" for key, value in doc.metadata.items() if key != "embedding"])
+        formatted_doc = f"- {doc.page_content}\n\n**Metadata:**\n{metadata_str}"
+        formatted_docs.append(formatted_doc)   
+    return "\n\n".join(formatted_docs)
+
+def stream_data(response):
+  for word in response.split(" "):
+        yield word + " "
+        time.sleep(0.05)
 
 if prompt := st.chat_input():
-    st.chat_message("human").write(prompt)
+    with st.chat_message("Human"):
+        st.markdown(prompt)
 
     # As usual, new messages are added to StreamlitChatMessageHistory when the Chain is called.
     config = {"configurable": {"session_id": "any"}}
-    response = conversational_rag_chain.invoke({"input": prompt}, config)
-    st.chat_message("ai").write(response["answer"])
+    res = get_response(prompt, config)
+
+    with st.chat_message("AI"):
+      st.write_stream(stream_data(res['answer']))
+      with st.popover("View Source"):
+        st.markdown("### Source Alerts 📢")
+        st.markdown(format_docs_with_metadata(res['context']))

requirements.txt

@@ -5,4 +5,5 @@ dnspython
 langchain 
 langchain-groq
 motor
-streamlit
+streamlit
+nest-asyncio