Update app.py

app.py

@@ -1,52 +1,82 @@
 import streamlit as st
+from huggingface_hub import InferenceClient
 import google.generativeai as genai
 
 # -------------------
-# API Key Setup
+# API Keys Setup
 # -------------------
+huggingface_token = st.secrets.get("HUGGINGFACE_HUB_TOKEN", "")
 gemini_api_key = st.secrets.get("GEN_API_KEY", "")
 
 # -------------------
-# Page Configuration
+# Configuration
 # -------------------
 st.set_page_config(page_title="Your AI Buddy", layout="wide")
 st.title("💡 Need answers? Just type below!")
 
+# List of recommended Hugging Face models (small ones to avoid storage issues)
+HF_RECOMMENDED_MODELS = ["google/flan-t5-large"]
+
 # -------------------
-# Gemini Setup
+# Sidebar Settings
 # -------------------
-if not gemini_api_key:
-    st.error("⚠️ Please set your 'GEN_API_KEY' in Streamlit secrets.")
-    st.stop()
+st.sidebar.title("⚙️ Settings")
+provider = st.sidebar.selectbox("Provider", ["Hugging Face", "Gemini"])
 
-genai.configure(api_key=gemini_api_key)
+# -------------------
+# Provider Setup
+# -------------------
+client = None
+model = None
 
-# Fetch available Gemini models
-available_models = [
-    m.name for m in genai.list_models()
-    if "generateContent" in m.supported_generation_methods
-]
+if provider == "Hugging Face":
+    if not huggingface_token:
+        st.error("⚠️ Please set your 'HUGGINGFACE_HUB_TOKEN' in Streamlit secrets.")
+        st.stop()
 
-if not available_models:
-    st.error("⚠️ No Gemini models available for your API key.")
-    st.stop()
+    selected_models = st.sidebar.multiselect(
+        "Choose HF models",
+        HF_RECOMMENDED_MODELS,
+        default=[HF_RECOMMENDED_MODELS[0]]
+    )
+    if not selected_models:
+        st.warning("⚠️ Please select at least one Hugging Face model.")
+        st.stop()
 
-# Reset session if old model is invalid
-if "model" in st.session_state and st.session_state["model"] not in available_models:
-    del st.session_state["model"]
+    client = InferenceClient(token=huggingface_token)
 
-model = st.sidebar.selectbox("Model", available_models, index=0)
+elif provider == "Gemini":
+    if not gemini_api_key:
+        st.error("⚠️ Please set your 'GEN_API_KEY' in Streamlit secrets.")
+        st.stop()
 
-# Initialize Gemini chat if needed
-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}")
+    genai.configure(api_key=gemini_api_key)
+
+    # Fetch available models that support generateContent
+    available_models = [
+        m.name for m in genai.list_models()
+        if "generateContent" in m.supported_generation_methods
+    ]
+    if not available_models:
+        st.error("⚠️ No Gemini models available for your API key.")
         st.stop()
 
+    # Reset session if old model is invalid
+    if "model" in st.session_state and st.session_state["model"] not in available_models:
+        del st.session_state["model"]
+
+    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
 # -------------------
@@ -64,8 +94,9 @@ if "messages" not in st.session_state:
 # Reset conversation button
 if st.sidebar.button("Reset Conversation"):
     st.session_state.messages = []
-    gemini_model = genai.GenerativeModel(model)
-    st.session_state.gemini_chat = gemini_model.start_chat(history=[])
+    if provider == "Gemini" and model:
+        gemini_model = genai.GenerativeModel(model)
+        st.session_state.gemini_chat = gemini_model.start_chat(history=[])
     st.rerun()
 
 # -------------------
@@ -79,21 +110,78 @@ for msg in st.session_state.messages:
 # User Input
 # -------------------
 if user_input := st.chat_input("Type your message..."):
-    # Show user message
     st.chat_message("user").markdown(user_input)
     st.session_state.messages.append({"role": "user", "content": user_input})
 
-    try:
-        with st.spinner("Gemini is thinking..."):
-            # Prepend system prompt for context
-            full_input = f"{system_prompt}\n\nUser: {user_input}"
-            resp = st.session_state.gemini_chat.send_message(full_input)
-            bot_text = resp.text
-    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()
+    # -------------------
+    # 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:
+                    # Prepare conversation history
+                    conv = "\n".join([
+                        f"{msg['role']}: {msg['content']}"
+                        for msg in st.session_state.messages
+                    ])
+                    prompt_text = f"{system_prompt}\n\n{conv}\nassistant:"
+
+                    # Always use text_generation (works for GPT, Bloom, Flan-T5)
+                    resp = client.text_generation(
+                        model=m,
+                        prompt=prompt_text,
+                        max_new_tokens=256,
+                        temperature=0.7,
+                        stop_sequences=["assistant:", "user:"],
+                        stream=False
+                    )
+
+                    # Parse response
+                    if isinstance(resp, str):
+                        bot_text = resp
+                    elif isinstance(resp, dict):
+                        bot_text = resp.get("generated_text", str(resp))
+                    elif isinstance(resp, list) and "generated_text" in resp[0]:
+                        bot_text = resp[0]["generated_text"]
+                    else:
+                        bot_text = "⚠️ Unexpected response format."
+
+                    # Strip if model repeats the prompt
+                    if bot_text.startswith(prompt_text):
+                        bot_text = bot_text[len(prompt_text):].strip()
+
+                except Exception as e:
+                    bot_text = f"⚠️ Error with **{m}**: {type(e).__name__}: {e}"
+
+                # Display final response
+                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()