Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
-from huggingface_hub import InferenceClient
-import google.generativeai as genai
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from langchain_huggingface import HuggingFaceEndpoint
+from langchain_google_genai import ChatGoogleGenerativeAI
 
 # -------------------
 # API Keys Setup
@@ -11,11 +12,14 @@ gemini_api_key = st.secrets.get("GEN_API_KEY", "")
 # -------------------
 # Configuration
 # -------------------
-st.set_page_config(page_title="Your AI Buddy", layout="wide")
-st.title("💡 Need answers? Just type below!")
+st.set_page_config(page_title="Your AI Buddy (LangChain)", layout="wide")
+st.title("🤖 Your AI Buddy (LangChain Edition)")
 
-# List of recommended Hugging Face models (small ones to avoid storage issues)
-HF_RECOMMENDED_MODELS = ["google/flan-t5-large"]
+HF_RECOMMENDED_MODELS = [
+    "google/flan-t5-small",
+    "google/flan-t5-base",
+    "distilgpt2"
+]
 
 # -------------------
 # Sidebar Settings
@@ -23,164 +27,94 @@ HF_RECOMMENDED_MODELS = ["google/flan-t5-large"]
 st.sidebar.title("⚙️ Settings")
 provider = st.sidebar.selectbox("Provider", ["Hugging Face", "Gemini"])
 
+system_prompt = st.sidebar.text_area(
+    "System Prompt",
+    "You are a helpful AI assistant. Provide concise and accurate answers."
+)
+
+# -------------------
+# Chat History
+# -------------------
+if "messages" not in st.session_state:
+    st.session_state.messages = [SystemMessage(content=system_prompt)]
+
+if st.sidebar.button("Reset Conversation"):
+    st.session_state.messages = [SystemMessage(content=system_prompt)]
+    st.rerun()
+
 # -------------------
 # Provider Setup
 # -------------------
-client = None
-model = None
+llm = None
 
 if provider == "Hugging Face":
     if not huggingface_token:
-        st.error("⚠️ Please set your 'HUGGINGFACE_HUB_TOKEN' in Streamlit secrets.")
+        st.error("⚠️ Missing Hugging Face token in secrets.")
         st.stop()
 
     selected_models = st.sidebar.multiselect(
-        "Choose HF models",
+        "Choose Hugging Face models",
         HF_RECOMMENDED_MODELS,
         default=[HF_RECOMMENDED_MODELS[0]]
     )
+
     if not selected_models:
-        st.warning("⚠️ Please select at least one Hugging Face model.")
+        st.warning("⚠️ Please select at least one model.")
         st.stop()
 
-    client = InferenceClient(token=huggingface_token)
+    # We'll just use the first one for now
+    model = selected_models[0]
+
+    llm = HuggingFaceEndpoint(
+        repo_id=model,
+        huggingfacehub_api_token=huggingface_token,
+        temperature=0.7,
+        max_new_tokens=256
+    )
 
 elif provider == "Gemini":
     if not gemini_api_key:
-        st.error("⚠️ Please set your 'GEN_API_KEY' in Streamlit secrets.")
+        st.error("⚠️ Missing Gemini API key in secrets.")
         st.stop()
 
-    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
-# -------------------
-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 = []
+    model = st.sidebar.selectbox(
+        "Choose Gemini model",
+        ["gemini-1.5-flash", "gemini-1.5-pro"],  # safe defaults
+        index=0
+    )
 
-# 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()
+    llm = ChatGoogleGenerativeAI(
+        model=model,
+        google_api_key=gemini_api_key,
+        temperature=0.7
+    )
 
 # -------------------
-# Display Chat Messages
+# Display History
 # -------------------
 for msg in st.session_state.messages:
-    with st.chat_message(msg["role"]):
-        st.markdown(msg["content"])
+    role = "user" if isinstance(msg, HumanMessage) else "assistant" if isinstance(msg, AIMessage) else "system"
+    if role != "system":  # hide system messages
+        with st.chat_message(role):
+            st.markdown(msg.content)
 
 # -------------------
 # 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...")
+    st.session_state.messages.append(HumanMessage(content=user_input))
 
+    # Generate response
+    with st.chat_message("assistant"):
+        with st.spinner("Thinking..."):
             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."
-
+                response = llm.invoke(st.session_state.messages)
+                bot_text = response.content if hasattr(response, "content") else str(response)
             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()
+                bot_text = f"⚠️ Error: {type(e).__name__}: {e}"
+
+        st.markdown(bot_text)
+        st.session_state.messages.append(AIMessage(content=bot_text))
+    st.rerun()