Update app.py

app.py

@@ -1,22 +1,22 @@
 import streamlit as st
-import whisper
-import numpy as np
+from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration
+from typing import List
 from langchain_community.llms import HuggingFaceEndpoint
 from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationalRetrievalChain
 from langchain_community.vectorstores import FAISS
-from langchain_core.prompts import PromptTemplate
-from langchain_core.memory import BaseMemory
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import RunnablePassthrough
-from langchain.chains import ConversationChain
+from langchain.prompts import PromptTemplate
 import os
 from dotenv import load_dotenv
 import requests
 from requests.adapters import HTTPAdapter
 from requests.packages.urllib3.util.retry import Retry
-import time
-from streamlit_chat import message
-from streamlit_audiorecorder import audiorecorder
+import whisper
+import numpy as np
+import av
+import time  # Added import time
+import queue
 
 # Load environment variables
 load_dotenv()
@@ -24,12 +24,15 @@ load_dotenv()
 # Initialize session state
 if "messages" not in st.session_state:
     st.session_state.messages = []
-if "audio_data" not in st.session_state:
-    st.session_state.audio_data = None
-if "recording" not in st.session_state:
+
+if "audio_buffer" not in st.session_state:
+    st.session_state.audio_buffer = queue.Queue()
+
+if 'recording' not in st.session_state:
     st.session_state.recording = False
-if "text_input" not in st.session_state:
-    st.session_state.text_input = ""
+
+if 'webrtc_ctx' not in st.session_state:
+    st.session_state.webrtc_ctx = None
 
 # Prompt template
 PROMPT_TEMPLATE = """
@@ -76,7 +79,7 @@ llm = HuggingFaceEndpoint(
 )
 
 # Setup memory and conversation chain
-memory = ConversationChain(
+memory = ConversationBufferMemory(
     memory_key="chat_history",
     return_messages=True
 )
@@ -95,24 +98,19 @@ qa_prompt = PromptTemplate(
     input_variables=["context", "chat_history", "question"]
 )
 
-def create_chain():
-    prompt = PromptTemplate(
-        template=PROMPT_TEMPLATE,
-        input_variables=["context", "chat_history", "question"]
-    )
-    
-    retriever = vectorstore.as_retriever()
-    
-    chain = (
-        {"context": retriever, "question": RunnablePassthrough()}
-        | prompt
-        | llm
-        | StrOutputParser()
-    )
-    
-    return chain
-
-conversation_chain = create_chain()
+conversation_chain = ConversationalRetrievalChain.from_llm(
+    llm=llm,
+    retriever=vectorstore.as_retriever(),
+    memory=memory,
+    combine_docs_chain_kwargs={"prompt": qa_prompt},
+    return_source_documents=True,
+    output_key='answer'  # Specify output_key to fix the error
+)
+
+def audio_frame_callback(frame: av.AudioFrame) -> av.AudioFrame:
+    audio = frame.to_ndarray().flatten()
+    st.session_state.audio_buffer.put(audio)
+    return frame
 
 def get_ai_response(user_input: str) -> str:
     max_retries = 3
@@ -128,13 +126,13 @@ def get_ai_response(user_input: str) -> str:
                 "question": user_input,
                 "chat_history": memory.chat_memory.messages[-5:]
             })
-            
+
             if not response or 'answer' not in response:
                 if attempt < max_retries - 1:
                     time.sleep(2 ** attempt)
                     continue
                 return "عذراً، كاين مشكل. حاول مرة أخرى."
-            
+
             return response['answer']
 
         except requests.exceptions.HTTPError as e:
@@ -142,7 +140,7 @@ def get_ai_response(user_input: str) -> str:
                 time.sleep(2 ** attempt)
                 continue
             return "عذراً، كاين مشكل مع النموذج. جرب سؤال أقصر."
-        
+
         except Exception as e:
             st.error(f"Error: {str(e)}")
             if attempt < max_retries - 1:
@@ -152,7 +150,7 @@ def get_ai_response(user_input: str) -> str:
 
 def process_message(user_input: str) -> None:
     st.session_state.messages.append({"role": "user", "content": user_input})
-    
+
     with st.spinner("جاري التفكير..."):
         ai_response = get_ai_response(user_input)
         if ai_response:
@@ -160,54 +158,70 @@ def process_message(user_input: str) -> None:
 
 def main():
     st.set_page_config(page_title="Darija AI Therapist", page_icon="🧠")
-    
+
     st.title("Darija AI Therapist 🧠")
     st.subheader("تكلم معايا بالدارجة على اللي كيجول فبالك")
-    
-    # Chat interface
-    # Create columns for text input and mic button
+
     col1, col2 = st.columns([9, 1])
     with col1:
         user_input = st.text_input("اكتب رسالتك هنا:", key="text_input")
     with col2:
-        # Mic button
         if st.session_state.recording:
-            mic_label = "🛑"
+            mic_icon = "🛑"
         else:
-            mic_label = "🎤"
-        
-        if st.button(mic_label):
+            mic_icon = "🎤"
+
+        if st.button(mic_icon):
             st.session_state.recording = not st.session_state.recording
             if st.session_state.recording:
-                st.session_state.audio_data = audiorecorder("Click to stop recording")
+                st.session_state.audio_buffer = queue.Queue()
+                st.session_state.webrtc_ctx = webrtc_streamer(
+                    key="speech-to-text",
+                    mode=WebRtcMode.SENDONLY,
+                    audio_receiver_size=256,
+                    rtc_configuration=RTCConfiguration(
+                        {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
+                    ),
+                    media_stream_constraints={"video": False, "audio": True},
+                    async_processing=True,
+                    audio_frame_callback=audio_frame_callback,
+                )
             else:
-                audio_data = st.session_state.audio_data
-                if audio_data is not None:
-                    # Convert byte data to numpy array
-                    audio_array = np.frombuffer(audio_data.tobytes(), dtype=np.int16)
-                    # Normalize audio data
-                    audio_array = audio_array.astype(np.float32) / np.iinfo(np.int16).max
-                    # Transcribe audio using Whisper
-                    result = whisper_model.transcribe(audio_array, language="ar")
-                    if result["text"]:
-                        # Put transcribed text into input field
-                        st.session_state.text_input = result["text"]
+                st.info("🔄 Processing audio...")
+                audio_frames = []
+                while not st.session_state.audio_buffer.empty():
+                    audio_frames.append(st.session_state.audio_buffer.get())
+
+                if audio_frames:
+                    audio_data = np.concatenate(audio_frames, axis=0).flatten()
+                    # Convert to 16-bit integers
+                    audio_data_int16 = (audio_data * 32767).astype(np.int16)
+                    # Use Whisper to transcribe
+                    result = whisper_model.transcribe(audio_data_int16, fp16=False)
+                    text = result.get("text", "")
+                    if text:
+                        process_message(text)
+                    else:
+                        st.warning("ما فهمتش الصوت. حاول مرة أخرى.")
+                    st.session_state.audio_buffer = queue.Queue()
                 else:
-                    st.error("No audio data recorded.")
+                    st.warning("ما تسجلش الصوت. حاول مرة أخرى.")
+                if st.session_state.webrtc_ctx:
+                    st.session_state.webrtc_ctx.stop()
+                    st.session_state.webrtc_ctx = None
+
+    if st.session_state.recording:
+        st.info("🎙️ Recording...")
+    else:
+        st.empty()
 
-    # Handle text submission
     if user_input:
         process_message(user_input)
-        st.session_state.text_input = ""  # Clear input field after sending
 
     # Display chat history
-    for message_data in st.session_state.messages:
-        role = message_data["role"]
-        content = message_data["content"]
-        if role == "user":
-            message(content, is_user=True)
-        else:
-            message(content)
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.write(message["content"])
 
 if __name__ == "__main__":
     main()