Update app.py

app.py

@@ -1,84 +1,52 @@
 import streamlit as st
-from huggingface_hub import InferenceClient
 import google.generativeai as genai
 
 # -------------------
-# API Keys Setup
+# API Key Setup
 # -------------------
-huggingface_token = st.secrets.get("HUGGINGFACE_HUB_TOKEN", "")
 gemini_api_key = st.secrets.get("GEN_API_KEY", "")
 
 # -------------------
-# Configuration
+# Page 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"]
-
-# -------------------
-# Sidebar Settings
 # -------------------
-st.sidebar.title("⚙️ Settings")
-provider = st.sidebar.selectbox("Provider", ["Hugging Face", "Gemini"])
-
+# Gemini Setup
 # -------------------
-# Provider Setup
-# -------------------
-client = None
-model = None
-
-if provider == "Hugging Face":
-    if not huggingface_token:
-        st.error("⚠️ Please set your 'HUGGINGFACE_HUB_TOKEN' in Streamlit secrets.")
-        st.stop()
+if not gemini_api_key:
+    st.error("⚠️ Please set your 'GEN_API_KEY' in Streamlit secrets.")
+    st.stop()
 
-    selected_models = st.sidebar.multiselect(
-        "Choose HF models",
-        HF_RECOMMENDED_MODELS,
-        default=[HF_RECOMMENDED_MODELS[0]]
-    )
+genai.configure(api_key=gemini_api_key)
 
-    if not selected_models:
-        st.warning("⚠️ Please select at least one Hugging Face model.")
-        st.stop()
+# Fetch available Gemini models
+available_models = [
+    m.name for m in genai.list_models()
+    if "generateContent" in m.supported_generation_methods
+]
 
-    client = InferenceClient(token=huggingface_token)
+if not available_models:
+    st.error("⚠️ No Gemini models available for your API key.")
+    st.stop()
 
-elif provider == "Gemini":
-    if not gemini_api_key:
-        st.error("⚠️ Please set your 'GEN_API_KEY' in Streamlit secrets.")
-        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"]
 
-    genai.configure(api_key=gemini_api_key)
+model = st.sidebar.selectbox("Model", available_models, index=0)
 
-    # 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.")
+# 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}")
         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
 # -------------------
@@ -96,9 +64,8 @@ if "messages" not in st.session_state:
 # 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=[])
+    gemini_model = genai.GenerativeModel(model)
+    st.session_state.gemini_chat = gemini_model.start_chat(history=[])
     st.rerun()
 
 # -------------------
@@ -112,73 +79,21 @@ 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})
 
-    # -------------------
-    # 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()
+    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()