Update app.py

app.py

@@ -70,3 +70,122 @@ elif provider == "Gemini":
     model = st.sidebar.selectbox("Model", available_models, index=0)
 
     # Initialize Gemini chat if model changes or not initialized
+    if "gemini_chat" not in st.session_state or st.session_state.get("model") != model:
+        st.session_state.model = model
+        try:
+            gemini_model = genai.GenerativeModel(model)
+            st.session_state.gemini_chat = gemini_model.start_chat(history=[])
+        except Exception as e:
+            st.error(f"⚠️ Could not initialize Gemini model: {e}")
+            st.stop()
+
+# -------------------
+# System Prompt
+# -------------------
+system_prompt = st.sidebar.text_area(
+    "System Prompt",
+    "You are a helpful AI assistant. Provide concise and accurate answers."
+)
+
+# -------------------
+# Chat History State
+# -------------------
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Reset conversation button
+if st.sidebar.button("Reset Conversation"):
+    st.session_state.messages = []
+    if provider == "Gemini" and model:
+        gemini_model = genai.GenerativeModel(model)
+        st.session_state.gemini_chat = gemini_model.start_chat(history=[])
+    st.rerun()
+
+# -------------------
+# Display Chat Messages
+# -------------------
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
+
+# -------------------
+# User Input
+# -------------------
+if user_input := st.chat_input("Type your message..."):
+    st.chat_message("user").markdown(user_input)
+    st.session_state.messages.append({"role": "user", "content": user_input})
+
+    # -------------------
+    # Hugging Face Logic
+    # -------------------
+    if provider == "Hugging Face":
+        for m in selected_models:
+            with st.chat_message("assistant"):
+                message_placeholder = st.empty()
+                message_placeholder.markdown(f"**{m}** is generating...")
+
+                try:
+                    conv = "\n".join(
+                        [f"{msg['role']}: {msg['content']}" for msg in st.session_state.messages]
+                    )
+
+                    if "flan-t5" in m or "t5" in m:
+                        # ✅ Use text2text for T5 models
+                        prompt_text = f"{system_prompt}\n\nUser asked:\n{conv}\n\nAnswer:"
+                        resp = client.text2text_generation(
+                            model=m,
+                            inputs=prompt_text,
+                            max_new_tokens=256,
+                            temperature=0.7,
+                        )
+                    else:
+                        # ✅ Use text_generation for GPT-style models
+                        stop_sequences = ["assistant:", "user:"]
+                        prompt_text = f"{system_prompt}\n\n{conv}\nassistant:"
+                        resp = client.text_generation(
+                            model=m,
+                            prompt=prompt_text,
+                            max_new_tokens=256,
+                            temperature=0.7,
+                            stop_sequences=stop_sequences,
+                            stream=False,
+                        )
+
+                    # Parse response safely
+                    if isinstance(resp, str):
+                        bot_text = resp
+                    elif isinstance(resp, dict) and "generated_text" in resp:
+                        bot_text = resp["generated_text"]
+                    elif isinstance(resp, list) and resp and "generated_text" in resp[0]:
+                        bot_text = resp[0]["generated_text"]
+                    else:
+                        bot_text = "⚠️ Model returned an unexpected format."
+
+                except Exception as e:
+                    bot_text = f"⚠️ Error with **{m}**: {type(e).__name__}: {e}"
+
+                final_response = f"**{m}**\n\n{bot_text}"
+                message_placeholder.markdown(final_response)
+                st.session_state.messages.append({"role": "assistant", "content": final_response})
+
+        st.rerun()
+
+    # -------------------
+    # Gemini Logic
+    # -------------------
+    elif provider == "Gemini":
+        try:
+            if user_input.strip():
+                with st.spinner("Gemini is thinking..."):
+                    resp = st.session_state.gemini_chat.send_message(user_input)
+                    bot_text = resp.text
+            else:
+                bot_text = "⚠️ Please enter a message before sending."
+        except Exception as e:
+            bot_text = f"⚠️ Gemini could not respond right now. Please try again. ({e})"
+
+        with st.chat_message("assistant"):
+            st.markdown(bot_text)
+
+        st.session_state.messages.append({"role": "assistant", "content": bot_text})
+        st.rerun()