Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,14 +7,17 @@ from huggingface_hub import InferenceClient
 # ==================================================
 # LOGGING
 # ==================================================
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s | %(levelname)s | %(message)s"
+)
 logger = logging.getLogger("AI-Agent")
 
 # ==================================================
 # RAW DATA (IN-MEMORY, 600+ ROWS)
 # ==================================================
 @st.cache_data
-def generate_data(rows=600):
+def generate_data(rows: int = 600):
     random.seed(42)
     data = []
 
@@ -35,7 +38,7 @@ def generate_data(rows=600):
             )
         })
 
-    logger.info("Raw data generated: %d rows", len(data))
+    logger.info("Generated raw dataset with %d rows", len(data))
     return data
 
 DATA = generate_data()
@@ -44,7 +47,7 @@ DATA = generate_data()
 # TOOLS (FUNCTION CALLING)
 # ==================================================
 def tool_get_stats():
-    logger.info("Tool called: get_stats")
+    logger.info("Tool call → get_stats")
 
     prices = [x["price"] for x in DATA]
     days = [x["delivery_days"] for x in DATA]
@@ -56,7 +59,7 @@ def tool_get_stats():
     }
 
 def tool_query_product(product: str):
-    logger.info("Tool called: query_product(%s)", product)
+    logger.info("Tool call → query_product(%s)", product)
 
     return [
         {
@@ -73,7 +76,7 @@ def tool_query_product(product: str):
 
 def tool_create_support_ticket(text: str):
     ticket_id = str(uuid.uuid4())[:8]
-    logger.info("Support ticket created: %s", ticket_id)
+    logger.info("Tool call → create_support_ticket (%s)", ticket_id)
 
     return {
         "ticket_id": ticket_id,
@@ -85,7 +88,7 @@ def tool_create_support_ticket(text: str):
 # ==================================================
 # SAFETY
 # ==================================================
-def is_dangerous(text: str):
+def is_dangerous(text: str) -> bool:
     blocked = ["delete", "drop", "truncate", "remove"]
     return any(b in text.lower() for b in blocked)
 
@@ -97,85 +100,95 @@ def agent(user_input: str, client: InferenceClient):
 
     if is_dangerous(user_input):
         logger.warning("Blocked dangerous operation")
-        return "❌ Dangerous operation is not allowed."
+        return "❌ Dangerous operations are not allowed."
 
     text = user_input.lower()
 
-    # ---- FUNCTION CALLS ----
+    # -------- FUNCTION CALLS --------
     if "stats" in text or "summary" in text:
-        s = tool_get_stats()
+        stats = tool_get_stats()
         return (
             f"📊 **Business Overview**\n"
-            f"- Rows: {s['rows']}\n"
-            f"- Avg Price: ${s['avg_price']}\n"
-            f"- Avg Delivery Days: {s['avg_delivery_days']}"
+            f"- Rows: {stats['rows']}\n"
+            f"- Avg Price: ${stats['avg_price']}\n"
+            f"- Avg Delivery Days: {stats['avg_delivery_days']}"
         )
 
     if text.startswith("show"):
         product = user_input.replace("show", "").strip()
-        result = tool_query_product(product)
+        rows = tool_query_product(product)
 
-        if not result:
+        if not rows:
             return (
-                "No data found for this product.\n\n"
+                "No records found for this product.\n\n"
                 "Would you like me to create a support ticket?"
             )
-        return result
+        return rows
 
     if "support" in text or "ticket" in text:
-        t = tool_create_support_ticket(user_input)
+        ticket = tool_create_support_ticket(user_input)
         return (
             f"🎫 **Support Ticket Created**\n"
-            f"- ID: {t['ticket_id']}\n"
-            f"- System: {t['system']}\n"
-            f"- Status: {t['status']}"
+            f"- ID: {ticket['ticket_id']}\n"
+            f"- System: {ticket['system']}\n"
+            f"- Status: {ticket['status']}"
         )
 
-    # ---- LLM FALLBACK (NO RAW DATA PASSED) ----
-    prompt = f"""
-You are an AI procurement assistant.
-You do NOT have access to raw data.
-If the user asks something unclear, suggest contacting support.
-
-User: {user_input}
-Assistant:
-"""
-    logger.info("LLM called")
-    return client.text_generation(
-        prompt,
-        max_new_tokens=150,
+    # -------- LLM FALLBACK (CHAT MODE) --------
+    logger.info("LLM fallback → chat_completion")
+
+    response = client.chat_completion(
+        messages=[
+            {
+                "role": "system",
+                "content": (
+                    "You are an AI procurement assistant. "
+                    "You do NOT have access to raw data. "
+                    "Answer concisely and professionally. "
+                    "If unsure, suggest creating a support ticket."
+                )
+            },
+            {
+                "role": "user",
+                "content": user_input
+            }
+        ],
+        max_tokens=150,
         temperature=0.3
     )
 
+    return response.choices[0].message.content
+
 # ==================================================
 # STREAMLIT UI
 # ==================================================
 st.set_page_config(page_title="AI Procurement Agent", layout="wide")
-st.title("🤖 AI Procurement Agent (Single-file MVP)")
+st.title("🤖 AI Procurement Agent — Single File MVP")
 
-# ---- Sidebar: Business Info ----
+# -------- Sidebar: Business Info --------
 stats = tool_get_stats()
-st.sidebar.header("📊 Business Info")
+
+st.sidebar.header("📊 Business Information")
 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
-""")
+st.sidebar.code(
+    "show tomato\n"
+    "database stats\n"
+    "create support ticket"
+)
 
-# ---- HF Token ----
+# -------- HF TOKEN --------
 hf_token = st.sidebar.text_input(
     "Hugging Face Token",
     type="password",
-    help="https://huggingface.co/settings/tokens"
+    help="Create token at https://huggingface.co/settings/tokens"
 )
 
 if not hf_token:
-    st.warning("Please provide HF token")
+    st.warning("Please provide a Hugging Face token.")
     st.stop()
 
 client = InferenceClient(
@@ -183,7 +196,7 @@ client = InferenceClient(
     token=hf_token
 )
 
-# ---- Chat ----
+# -------- Chat --------
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
@@ -191,8 +204,8 @@ user_input = st.chat_input("Ask the procurement agent...")
 
 if user_input:
     st.session_state.messages.append(("user", user_input))
-    response = agent(user_input, client)
-    st.session_state.messages.append(("assistant", response))
+    reply = agent(user_input, client)
+    st.session_state.messages.append(("assistant", reply))
 
-for role, msg in st.session_state.messages:
-    st.chat_message(role).write(msg)
+for role, message in st.session_state.messages:
+    st.chat_message(role).write(message)