Update app.py

app.py

@@ -113,18 +113,9 @@ if "documents" in st.session_state:
         "Identify research gaps",
         "Suggest research ideas",
         "Simulate a debate",
-        "Translate response",
         "Generate citation"
     ])
 
-    # Translation language input with option to add custom
-    default_languages = ["Spanish", "French", "German", "Chinese", "Urdu", "Other"]
-    selected_language = st.selectbox("🌍 Choose translation language (only for Translate task):", default_languages)
-    if selected_language == "Other":
-        user_language = st.text_input("Please enter your desired language:", key="custom_lang")
-    else:
-        user_language = selected_language
-
     if st.button("🚀 Run Agent"):
         with st.spinner("Running agents..."):
             docs = st.session_state.documents[:10]
@@ -170,17 +161,6 @@ if "documents" in st.session_state:
                 st.markdown("### 🎭 Debate")
                 st.write(debate)
 
-            # Translate agent
-            elif task == "Translate response":
-                if "last_agent_output" in st.session_state:
-                    to_translate = st.session_state["last_agent_output"]
-                    translate_chain = LLMChain(llm=llm, prompt=translate_prompt)
-                    translated = translate_chain.invoke({"language": user_language, "content": to_translate})
-                    st.markdown(f"### 🌐 Translated Response ({user_language})")
-                    st.write(translated)
-                else:
-                    st.warning("Please run another agent first (e.g., Summary, Gap, Debate) to generate something to translate.")
-
             # Citation agent
             elif task == "Generate citation":
                 citation_chain = create_stuff_documents_chain(llm, citation_prompt)
@@ -188,3 +168,24 @@ if "documents" in st.session_state:
                 st.session_state["last_agent_output"] = citation
                 st.markdown("### 📎 APA Citation")
                 st.code(citation, language="markdown")
+
+        # Translation toggle after agent response
+        if "last_agent_output" in st.session_state:
+            if st.toggle("🌍 Translate the response?"):
+                default_languages = ["Spanish", "French", "German", "Chinese", "Urdu", "Other"]
+                selected_language = st.selectbox("Choose translation language:", default_languages)
+                if selected_language == "Other":
+                    user_language = st.text_input("Please enter your desired language:", key="custom_lang")
+                else:
+                    user_language = selected_language
+
+                to_translate = st.session_state["last_agent_output"]
+                if isinstance(to_translate, dict):
+                    combined_text = "\n\n".join(str(value) for value in to_translate.values())
+                else:
+                    combined_text = str(to_translate)
+
+                translate_chain = LLMChain(llm=llm, prompt=translate_prompt)
+                translated = translate_chain.invoke({"language": user_language, "content": combined_text})
+                st.markdown(f"### 🌐 Translated Response ({user_language})")
+                st.write(translated)