Update app.py

app.py

@@ -129,7 +129,7 @@ for message in st.session_state.messages:
 def on_select():
     st.session_state.disabled = True
 
-    
+
 def get_message_history():
     for message in st.session_state.messages:
         role, content = message["role"], message["content"]
@@ -147,11 +147,12 @@ if prompt := st.chat_input("How can I help you today?"):
         message_placeholder = st.empty()
         full_response = ""
         message_history = "\n".join(list(get_message_history())[-3:])
-        #question = st.text_input("Ask your question", placeholder="Try to include context in your question")
-        result = qa_chain(prompt)
+        question = st.text_input("Ask your question", placeholder="Try to include context in your question")
+        result = qa_chain(question)
         output = [result['result']]
 
     def generate_pdf():
+        generate_audio()
         page_number = int(result['source_documents'][0].metadata['page'])
         doc = fitz.open(str(result['source_documents'][0].metadata['source']))
         text = str(result['source_documents'][0].page_content)
@@ -171,9 +172,11 @@ if prompt := st.chat_input("How can I help you today?"):
             pix.save(output_image, "png")
             pdf_document.close()
         pdf_page_to_image('/home/user/app/pdf2image/output.pdf', page_number, '/home/user/app/pdf2image/output.png')
-        image = Image.open('/home/user/app/pdf2image/output.png')
-        st.image(image)
-        st.session_state.image_displayed = True
+        # image = Image.open('/home/user/app/pdf2image/output.png')
+        # st.image(image)
+        # st.session_state.image_displayed = True
+        path = '/home/user/app/pdf2image/output.png'
+        return path
 
     def generate_audio():
         sound_file = BytesIO()
@@ -188,8 +191,25 @@ if prompt := st.chat_input("How can I help you today?"):
         message_placeholder.markdown(full_response + "▌")
         message_placeholder.markdown(full_response)
 
-    if st.toggle("Reference"):
-        generate_pdf()
+    def display_image(image_path):
+        """Displays an image at the given path."""
+        st.image(image_path)
+
+
+    # Create a button to display the image.
+    if st.button("Display Image"):
+        # Get the image path from the session state dictionary.
+        image_path = st.session_state.get("image_path")
+
+        # If the image path is not set, set it to the default image path.
+        if image_path is None:
+            image_path = generate_pdf()
+
+        # Display the image.
+        display_image(image_path)
+
+    # Store the image path in the session state dictionary.
+    st.session_state["image_path"] = image_path
 
     st.session_state.messages.append({"role": "assistant", "content": full_response})