Update app.py

app.py

@@ -78,21 +78,25 @@ def translate_text(text, dest_language="hi"):
 
 @st.cache_resource
 def get_vectorstore():
-    embedding_model = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
-    os.makedirs("vectorstore/db_faiss", exist_ok=True)
-    
-    faiss_files = ["index.faiss", "index.pkl"]
-    for filename in faiss_files:
-        if not os.path.exists(f"vectorstore/db_faiss/{filename}"):
-            hf_hub_download(
-                repo_id=VECTORSTORE_REPO_ID,
-                filename=filename,
-                local_dir="vectorstore/db_faiss",
-                token=HF_TOKEN,
-                repo_type="dataset"
-            )
-    
-    return FAISS.load_local("vectorstore/db_faiss", embedding_model, allow_dangerous_deserialization=True)
+    try:
+        embedding_model = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
+        os.makedirs("vectorstore/db_faiss", exist_ok=True)
+        
+        faiss_files = ["index.faiss", "index.pkl"]
+        for filename in faiss_files:
+            if not os.path.exists(f"vectorstore/db_faiss/{filename}"):
+                hf_hub_download(
+                    repo_id=VECTORSTORE_REPO_ID,
+                    filename=filename,
+                    local_dir="vectorstore/db_faiss",
+                    token=HF_TOKEN,
+                    repo_type="dataset"
+                )
+        
+        return FAISS.load_local("vectorstore/db_faiss", embedding_model, allow_dangerous_deserialization=True)
+    except Exception as e:
+        st.error(f"Vectorstore initialization failed: {str(e)}")
+        st.stop()
 
 def set_custom_prompt(custom_prompt_template):
     return PromptTemplate(template=custom_prompt_template, input_variables=["context", "question"])
@@ -144,23 +148,35 @@ def handle_user_input(prompt, qa_chain):
             st.error(f"Error: {str(e)}")
 
 def handle_translation():
-    if "last_response" in st.session_state and st.session_state.last_response:
+    if st.session_state.get("last_response") and not st.session_state.get("translation_done", False):
         try:
-            if not st.session_state.get("translation_done", False):
-                translated_text = translate_text(st.session_state.last_response, "hi")
-                
-                # Update the existing assistant message
-                for msg in reversed(st.session_state.messages):
-                    if msg["role"] == "assistant":
-                        msg["content"] = f'<div class="hindi-text">{translated_text}</div>'
-                        msg["animated"] = False  # Reset animation state
-                        break
-                
-                st.session_state.translation_done = True
-                st.rerun()
-                
+            translated_text = translate_text(st.session_state.last_response, "hi")
+            
+            # Update all assistant messages (in case of multiple responses)
+            for idx, msg in enumerate(st.session_state.messages):
+                if msg["role"] == "assistant":
+                    # Preserve original English text in metadata
+                    st.session_state.messages[idx] = {
+                        "role": "assistant",
+                        "content": f'<div class="hindi-text">{translated_text}</div>',
+                        "original": st.session_state.messages[idx]["content"]
+                    }
+            
+            st.session_state.translation_done = True
+            st.experimental_rerun()  # Force full refresh
+            
         except Exception as e:
             st.error(f"Translation error: {str(e)}")
+            st.session_state.translation_done = False
+
+def render_chat_messages():
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"], avatar="🐿" if message["role"] == "user" else "🪈"):
+            content = message.get("original", message["content"])  # Show original if available
+            if "hindi-text" in message["content"]:
+                st.markdown(message["content"], unsafe_allow_html=True)
+            else:
+                st.markdown(content)
 
 def main():
     st.markdown("""
@@ -297,8 +313,13 @@ def main():
             handle_user_input(prompt, qa_chain)
 
         if st.session_state.get("last_response"):
-            if st.button("🌐 Translate to Hindi", key="translate_btn"):
-                handle_translation()
+            col1, col2 = st.columns([1, 3])
+            with col1:
+                if st.button("🌐 Translate to Hindi", key="translate_btn"):
+                    handle_translation()
+            with col2:
+                if st.session_state.get("translation_done"):
+                    st.success("Translation to Hindi completed!")
 
     except Exception as e:
         st.error(f"Initialization error: {str(e)}")