Update src/streamlit_app.py

src/streamlit_app.py

@@ -16,10 +16,6 @@ logger = logging.getLogger("HF-AI-Agent")
 # CONFIG
 # ======================
 HF_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
-HF_TOKEN = os.getenv("HF_TOKEN")  # HF Spaces Secret
-
-if HF_TOKEN is None:
-    raise RuntimeError("HF_TOKEN secret is missing")
 
 # ======================
 # LOAD DATA (600+ rows)
@@ -49,7 +45,7 @@ def load_data():
 df = load_data()
 
 # ======================
-# TOOLS (FUNCTION CALLS)
+# TOOLS
 # ======================
 def query_database(product: str):
     logger.info("Tool call: query_database(%s)", product)
@@ -74,17 +70,44 @@ def create_support_ticket(text: str):
     }
 
 # ======================
-# SAFETY GUARD
+# SAFETY
 # ======================
 def is_dangerous(text: str):
     blocked = ["delete", "drop", "truncate", "remove table"]
     return any(b in text.lower() for b in blocked)
 
 # ======================
-# AGENT LOGIC
+# STREAMLIT UI
+# ======================
+st.set_page_config(page_title="AI Procurement Agent (MVP)", layout="wide")
+st.title("🤖 AI Procurement Agent — MVP")
+
+st.caption("Hugging Face powered, data-aware, safe AI agent demo")
+
+# ======================
+# TOKEN HANDLING (IMPORTANT FIX)
 # ======================
-client = InferenceClient(model=HF_MODEL, token=HF_TOKEN)
+hf_token = os.getenv("HF_TOKEN")
+
+if not hf_token:
+    hf_token = st.sidebar.text_input(
+        "🔑 Hugging Face API Token",
+        type="password",
+        help="Create a token at https://huggingface.co/settings/tokens"
+    )
+
+if not hf_token:
+    st.warning("Please provide a Hugging Face API token to continue.")
+    st.stop()
 
+client = InferenceClient(
+    model=HF_MODEL,
+    token=hf_token
+)
+
+# ======================
+# AGENT LOGIC
+# ======================
 def agent(user_input: str):
     if is_dangerous(user_input):
         logger.warning("Blocked unsafe operation")
@@ -124,17 +147,15 @@ User: {user_input}
 Assistant:
 """
     logger.info("HF model called")
-    return client.text_generation(prompt, max_new_tokens=200, temperature=0.3)
+    return client.text_generation(
+        prompt,
+        max_new_tokens=200,
+        temperature=0.3
+    )
 
 # ======================
-# STREAMLIT UI
+# SIDEBAR – BUSINESS INFO
 # ======================
-st.set_page_config(page_title="AI Procurement Agent (MVP)", layout="wide")
-st.title("🤖 AI Procurement Agent — MVP (HF Spaces)")
-
-st.caption("Minimal, secure, data-aware AI agent demo")
-
-# Sidebar (Business Info)
 st.sidebar.header("📊 Business Overview")
 stats = get_aggregates()
 st.sidebar.metric("Rows", stats["rows"])
@@ -148,7 +169,9 @@ database stats
 create support ticket
 """)
 
-# Chat
+# ======================
+# CHAT
+# ======================
 if "messages" not in st.session_state:
     st.session_state.messages = []