Update app.py

app.py

@@ -24,8 +24,7 @@ from PIL import Image
 from langchain.vectorstores import FAISS
 import transformers
 from pydub import AudioSegment
-#from streamlit_extras.stateful_button import button
-from streamlit_card import card
+from streamlit_extras.stateful_button import button
 
 user_session_id = uuid.uuid4()
 
@@ -149,10 +148,6 @@ if prompt := st.chat_input("How can I help you today?"):
         message_history = "\n".join(list(get_message_history())[-3:])
         result = qa_chain(prompt)
         output = [result['result']]
-        for item in output:
-            full_response += item
-            message_placeholder.markdown(full_response + "▌")
-            message_placeholder.markdown(full_response)
 
     def generate_pdf():
         page_number = int(result['source_documents'][0].metadata['page'])
@@ -165,7 +160,7 @@ if prompt := st.chat_input("How can I help you today?"):
                     highlight = page.add_highlight_annot(inst)
                     highlight.update()
         doc.save("/home/user/app/pdf2image/output.pdf", garbage=4, deflate=True, clean=True)
-    
+
         def pdf_page_to_image(pdf_file, page_number, output_image):
             pdf_document = fitz.open(pdf_file)
             page = pdf_document[page_number]
@@ -174,53 +169,38 @@ 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')
+        pdfPath = '/home/user/app/pdf2image/output.png'
+        return pdfPath
         #image = Image.open('/home/user/app/pdf2image/output.png')
-        st.image('/home/user/app/pdf2image/output.png')
         #message_placeholder.image(image)
         #st.session_state.reference = True
 
 
+
     def generate_audio():
         with open('/home/user/app/audio/audio.mp3','wb') as sound_file:
             tts = gTTS(result['result'], lang='en', tld='co.in')
             tts.write_to_fp(sound_file)
         sound = AudioSegment.from_mp3("/home/user/app/audio/audio.mp3")
         sound.export("/home/user/app/audio/audio.wav", format="wav")
-        st.audio("/home/user/app/audio/audio.wav")
-
-    card(title="Reference",
-         text='Ref',
-         url=None,
-         on_click=generate_pdf)
-
-    card(title='Audio',
-         text='Audioo',
-        url = None,
-        on_click = generate_audio)
-
-    # if "reference" not in st.session_state:
-    #     st.session_state.reference = '/home/user/app/pdf2image/default_output.png'
-    # if "audio" not in st.session_state:
-    #     st.session_state.audio = "/home/user/app/audio/beep-06.mp3"
-
-
-    # with st.sidebar:
-    #     choice = st.radio("References and TTS",["Reference" , 'TTS'], index=None,)
-
-    #     if choice == 'Reference':
-    #         generate_pdf()
-    #         st.session_state['reference'] = '/home/user/app/pdf2image/output.png'
-    #         st.image(st.session_state['reference'])
-
-    #     if choice == 'TTS':
-    #         with open('/home/user/app/audio/audio.mp3','wb') as sound_file:
-    #             tts = gTTS(result['result'], lang='en', tld = 'co.in')
-    #             tts.write_to_fp(sound_file)
-    #         sound = AudioSegment.from_mp3("/home/user/app/audio/audio.mp3")
-    #         sound.export("/home/user/app/audio/audio.wav", format="wav")
-    #         st.session_state['audio'] = '/home/user/app/audio/audio.wav'
-    #         st.audio(st.session_state['audio'])
-        
-    st.session_state.messages.append({"role": "assistant", "content": full_response})
-    
+        audioPath = "/home/user/app/audio/audio.wav"
+        return audioPath
     
+    for item in output:
+        full_response += item
+        message_placeholder.markdown(full_response + "▌")
+        message_placeholder.markdown(full_response)
+
+    if button('Audio', key='audio'):
+        ap = generate_audio()
+        st.audio(ap)
+
+    if button('Ref', key='ref'):
+        pp = generate_pdf()
+        st.image(pp)
+
+
+
+    st.session_state.messages.append({"role": "assistant", "content": full_response})
+
+