Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,7 @@
 import os
 import uuid
 import logging
+import json
 import pandas as pd
 import numpy as np
 import streamlit as st
@@ -18,13 +19,13 @@ logger = logging.getLogger("HF-AI-Agent")
 HF_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
 
 # ======================
-# LOAD DATA (600+ rows)
+# LOAD DATA
 # ======================
 @st.cache_data
 def load_data():
     np.random.seed(42)
     rows = 600
-    df = pd.DataFrame({
+    return pd.DataFrame({
         "order_id": range(1, rows + 1),
         "supplier": np.random.choice(
             ["Supplier A", "Supplier B", "Supplier C", "Supplier D"], rows
@@ -39,8 +40,6 @@ def load_data():
             ["Delivered", "Pending", "Delayed"], rows
         ),
     })
-    logger.info("Dataset loaded (%d rows)", len(df))
-    return df
 
 df = load_data()
 
@@ -48,11 +47,9 @@ df = load_data()
 # TOOLS
 # ======================
 def query_database(product: str):
-    logger.info("Tool call: query_database(%s)", product)
     return df[df["product"].str.lower() == product.lower()].head(5)
 
 def get_aggregates():
-    logger.info("Tool call: get_aggregates()")
     return {
         "rows": len(df),
         "avg_price": round(df["price"].mean(), 2),
@@ -60,10 +57,8 @@ def get_aggregates():
     }
 
 def create_support_ticket(text: str):
-    ticket_id = str(uuid.uuid4())[:8]
-    logger.info("Support ticket created: %s", ticket_id)
     return {
-        "ticket_id": ticket_id,
+        "ticket_id": str(uuid.uuid4())[:8],
         "system": "GitHub Issues (mock)",
         "status": "Created",
         "description": text,
@@ -77,60 +72,57 @@ def is_dangerous(text: str):
     return any(b in text.lower() for b in blocked)
 
 # ======================
-# STREAMLIT UI
+# JSON BUILDER (1-QADAM)
 # ======================
-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")
+def build_request_json(user_input: str):
+    intent = "llm"
 
-# ======================
-# TOKEN HANDLING (IMPORTANT FIX)
-# ======================
-hf_token = os.getenv("HF_TOKEN")
+    text = user_input.lower()
 
-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 "stats" in text or "summary" in text:
+        intent = "stats"
+    elif text.startswith("show"):
+        intent = "query"
+    elif "support" in text or "ticket" in text:
+        intent = "support"
 
-if not hf_token:
-    st.warning("Please provide a Hugging Face API token to continue.")
-    st.stop()
+    request_json = {
+        "request_id": str(uuid.uuid4()),
+        "intent": intent,
+        "payload": {
+            "text": user_input
+        }
+    }
 
-client = InferenceClient(
-    model=HF_MODEL,
-    token=hf_token
-)
+    logger.info("REQUEST JSON: %s", request_json)
+    return request_json
 
 # ======================
-# AGENT LOGIC
+# ROUTER (2-QADAM)
 # ======================
-def agent(user_input: str):
-    if is_dangerous(user_input):
-        logger.warning("Blocked unsafe operation")
+def router(req: dict, client: InferenceClient):
+    intent = req["intent"]
+    text = req["payload"]["text"]
+
+    if is_dangerous(text):
         return "❌ This operation is not allowed."
 
-    if "stats" in user_input or "summary" in user_input:
+    if intent == "stats":
         s = get_aggregates()
         return (
             f"📊 **Database Summary**\n"
             f"- Rows: {s['rows']}\n"
-            f"- Average Price: ${s['avg_price']}\n"
+            f"- Avg Price: ${s['avg_price']}\n"
             f"- Avg Delivery Days: {s['avg_delivery_days']}"
         )
 
-    if user_input.lower().startswith("show"):
-        product = user_input.replace("show", "").strip()
+    if intent == "query":
+        product = text.replace("show", "").strip()
         result = query_database(product)
-        if result.empty:
-            return "No records found."
-        return result
+        return "No records found." if result.empty else result
 
-    if "support" in user_input or "ticket" in user_input:
-        ticket = create_support_ticket(user_input)
+    if intent == "support":
+        ticket = create_support_ticket(text)
         return (
             f"🎫 **Support Ticket Created**\n"
             f"- ID: {ticket['ticket_id']}\n"
@@ -138,15 +130,15 @@ def agent(user_input: str):
             f"- Status: {ticket['status']}"
         )
 
+    # LLM fallback
     prompt = f"""
 You are an AI procurement assistant.
-You do NOT have direct database access.
-Answer concisely and professionally.
+You do NOT have database access.
+Answer shortly and professionally.
 
-User: {user_input}
+User: {text}
 Assistant:
 """
-    logger.info("HF model called")
     return client.text_generation(
         prompt,
         max_new_tokens=200,
@@ -154,24 +146,29 @@ Assistant:
     )
 
 # ======================
-# SIDEBAR – BUSINESS INFO
+# STREAMLIT UI
 # ======================
-st.sidebar.header("📊 Business Overview")
+st.set_page_config(page_title="AI Procurement Agent (JSON-based)", layout="wide")
+st.title("🤖 AI Procurement Agent — JSON Router MVP")
+
+hf_token = os.getenv("HF_TOKEN") or st.sidebar.text_input(
+    "🔑 Hugging Face API Token",
+    type="password"
+)
+
+if not hf_token:
+    st.warning("HF token required")
+    st.stop()
+
+client = InferenceClient(model=HF_MODEL, token=hf_token)
+
+# Sidebar
 stats = get_aggregates()
 st.sidebar.metric("Rows", stats["rows"])
 st.sidebar.metric("Avg Price", f"${stats['avg_price']}")
 st.sidebar.metric("Avg Delivery Days", stats["avg_delivery_days"])
 
-st.sidebar.markdown("### Sample Queries")
-st.sidebar.code("""
-show tomato
-database stats
-create support ticket
-""")
-
-# ======================
-# CHAT
-# ======================
+# Chat memory
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
@@ -179,11 +176,12 @@ user_input = st.chat_input("Ask the procurement agent...")
 
 if user_input:
     st.session_state.messages.append(("user", user_input))
-    response = agent(user_input)
+
+    request_json = build_request_json(user_input)
+    response = router(request_json, client)
+
     st.session_state.messages.append(("assistant", response))
 
+# Render chat
 for role, msg in st.session_state.messages:
-    if role == "user":
-        st.chat_message("user").write(msg)
-    else:
-        st.chat_message("assistant").write(msg)
+    st.chat_message(role).write(msg)