Update app.py

app.py

@@ -123,46 +123,32 @@ def create_conversational_rag_chain(vectorstore):
     return create_rag_chain()
 conversational_rag_chain = create_conversational_rag_chain()
 
-def display_chat_history(chat_history):
-    """Displays the chat history in Streamlit."""
-    for msg in chat_history.messages:
-        st.chat_message(msg.type).write(msg.content)
-
-
-
-def main_page(conversational_rag_chain):
-    load_dotenv()
-    """Main page for the Streamlit app."""
+def main(conversational_rag_chain):
+    """Main function for the Streamlit app."""
+    # Initialize chat history if not already present in session state
     msgs = st.session_state.get("chat_history", StreamlitChatMessageHistory(key="special_app_key"))
     chain_with_history = conversational_rag_chain
 
     st.title("Conversational RAG Chatbot")
 
-    display_chat_history(msgs)
-
     if prompt := st.chat_input():
         st.chat_message("human").write(prompt)
 
+        # Prepare the input dictionary with the correct keys
         input_dict = {"input": prompt, "chat_history": msgs.messages}
         config = {"configurable": {"session_id": "any"}}
 
+        # Process user input and handle response
         response = chain_with_history.invoke(input_dict, config)
         st.chat_message("ai").write(response["answer"])
 
+        # Display retrieved documents (if any and present in response)
         if "docs" in response and response["documents"]:
-            docs = response["documents"]
-            def expand_document(index):
                 st.write(f"Expanding document {index+1}...")
-            display_documents(docs, expand_document)
+            display_documents(docs, expand_document)  # Pass click function
 
+    # Update chat history in session state
     st.session_state["chat_history"] = msgs
 
-
-def main():
-    """Main function for the Streamlit app with page navigation."""
-
-    main_page(conversational_rag_chain)
- 
-
 if __name__ == "__main__":
-    main()
+    main(conversational_rag_chain)