Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,27 +2,66 @@ import streamlit as st
 from huggingface_hub import InferenceClient
 import os
 
+# ==============================
+# PAGE CONFIG
+# ==============================
 st.set_page_config(page_title="AI Assistant", layout="wide")
+st.title("🤖 AI Assistant")
 
-# Initialize client
-client = InferenceClient(
-    model="meta-llama/Meta-Llama-3-8B",
-    token=os.environ.get("HF_TOKEN")
-)
+# ==============================
+# LOAD MODEL CLIENT
+# ==============================
+@st.cache_resource
+def load_client():
+    return InferenceClient(
+        model="meta-llama/Llama-3-8b-instruct",
+        token=os.environ.get("HF_TOKEN")
+    )
 
-st.title("🤖 AI Assistant")
+client = load_client()
 
-query = st.text_input("Ask something")
+# ==============================
+# SESSION STATE (CHAT HISTORY)
+# ==============================
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
-if st.button("Generate"):
-    if not query:
-        st.warning("Please enter a query")
+# ==============================
+# DISPLAY CHAT HISTORY
+# ==============================
+for msg in st.session_state.messages:
+    if msg["role"] == "user":
+        st.chat_message("user").write(msg["content"])
     else:
-        with st.spinner("Generating response..."):
-            response = client.text_generation(
-                query,
-                max_new_tokens=150
+        st.chat_message("assistant").write(msg["content"])
+
+# ==============================
+# USER INPUT
+# ==============================
+query = st.chat_input("Ask anything...")
+
+if query:
+    # Store user message
+    st.session_state.messages.append({"role": "user", "content": query})
+    st.chat_message("user").write(query)
+
+    try:
+        with st.spinner("Thinking..."):
+
+            # ✅ Chat-based request (BEST PRACTICE)
+            response = client.chat_completion(
+                messages=[
+                    {"role": "system", "content": "You are a helpful, professional AI assistant."}
+                ] + st.session_state.messages,
+                max_tokens=300,
+                temperature=0.7,
             )
 
-        st.subheader("Response")
-        st.write(response)
+            reply = response.choices[0].message["content"]
+
+        # Store assistant reply
+        st.session_state.messages.append({"role": "assistant", "content": reply})
+        st.chat_message("assistant").write(reply)
+
+    except Exception as e:
+        st.error(f"Error: {str(e)}")