Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,51 +1,80 @@
-import sqlite3
+import random
 import uuid
 import logging
-import pandas as pd
 import streamlit as st
 from huggingface_hub import InferenceClient
 
-# ======================
+# ==================================================
 # LOGGING
-# ======================
+# ==================================================
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("AI-Agent")
 
-# ======================
-# DB CONNECTION
-# ======================
-@st.cache_resource
-def get_db():
-    return sqlite3.connect("procurement.db", check_same_thread=False)
-
-conn = get_db()
-
-# ======================
-# TOOLS (FUNCTION CALLS)
-# ======================
+# ==================================================
+# RAW DATA (IN-MEMORY, 600+ ROWS)
+# ==================================================
+@st.cache_data
+def generate_data(rows=600):
+    random.seed(42)
+    data = []
+
+    for i in range(1, rows + 1):
+        data.append({
+            "order_id": i,
+            "supplier": random.choice(
+                ["Supplier A", "Supplier B", "Supplier C", "Supplier D"]
+            ),
+            "product": random.choice(
+                ["Tomato", "Cheese", "Flour", "Oil", "Meat"]
+            ),
+            "quantity": random.randint(1, 100),
+            "price": round(random.uniform(5, 50), 2),
+            "delivery_days": random.randint(1, 14),
+            "status": random.choice(
+                ["Delivered", "Pending", "Delayed"]
+            )
+        })
+
+    logger.info("Raw data generated: %d rows", len(data))
+    return data
+
+DATA = generate_data()
+
+# ==================================================
+# TOOLS (FUNCTION CALLING)
+# ==================================================
 def tool_get_stats():
     logger.info("Tool called: get_stats")
-    q = """
-    SELECT COUNT(*) rows,
-           ROUND(AVG(price),2) avg_price,
-           ROUND(AVG(delivery_days),2) avg_delivery_days
-    FROM orders
-    """
-    return pd.read_sql(q, conn).iloc[0].to_dict()
-
-def tool_query_product(product):
+
+    prices = [x["price"] for x in DATA]
+    days = [x["delivery_days"] for x in DATA]
+
+    return {
+        "rows": len(DATA),
+        "avg_price": round(sum(prices) / len(prices), 2),
+        "avg_delivery_days": round(sum(days) / len(days), 2),
+    }
+
+def tool_query_product(product: str):
     logger.info("Tool called: query_product(%s)", product)
-    q = """
-    SELECT order_id, supplier, quantity, price, delivery_days, status
-    FROM orders
-    WHERE LOWER(product)=LOWER(?)
-    LIMIT 5
-    """
-    return pd.read_sql(q, conn, params=(product,))
-
-def tool_create_ticket(text):
+
+    return [
+        {
+            "order_id": x["order_id"],
+            "supplier": x["supplier"],
+            "quantity": x["quantity"],
+            "price": x["price"],
+            "delivery_days": x["delivery_days"],
+            "status": x["status"],
+        }
+        for x in DATA
+        if x["product"].lower() == product.lower()
+    ][:5]
+
+def tool_create_support_ticket(text: str):
     ticket_id = str(uuid.uuid4())[:8]
     logger.info("Support ticket created: %s", ticket_id)
+
     return {
         "ticket_id": ticket_id,
         "system": "GitHub Issues (mock)",
@@ -53,67 +82,80 @@ def tool_create_ticket(text):
         "description": text,
     }
 
-# ======================
+# ==================================================
 # SAFETY
-# ======================
-def is_dangerous(text):
+# ==================================================
+def is_dangerous(text: str):
     blocked = ["delete", "drop", "truncate", "remove"]
     return any(b in text.lower() for b in blocked)
 
-# ======================
+# ==================================================
 # AGENT (INTENT → FUNCTION CALL)
-# ======================
-def agent(user_input, client):
+# ==================================================
+def agent(user_input: str, client: InferenceClient):
+    logger.info("User input: %s", user_input)
+
     if is_dangerous(user_input):
-        logger.warning("Blocked dangerous request")
+        logger.warning("Blocked dangerous operation")
         return "❌ Dangerous operation is not allowed."
 
     text = user_input.lower()
 
+    # ---- FUNCTION CALLS ----
     if "stats" in text or "summary" in text:
-        stats = tool_get_stats()
+        s = tool_get_stats()
         return (
             f"📊 **Business Overview**\n"
-            f"- Rows: {stats['rows']}\n"
-            f"- Avg Price: ${stats['avg_price']}\n"
-            f"- Avg Delivery Days: {stats['avg_delivery_days']}"
+            f"- Rows: {s['rows']}\n"
+            f"- Avg Price: ${s['avg_price']}\n"
+            f"- Avg Delivery Days: {s['avg_delivery_days']}"
         )
 
     if text.startswith("show"):
         product = user_input.replace("show", "").strip()
-        df = tool_query_product(product)
-        if df.empty:
-            return "No data found. Would you like me to create a support ticket?"
-        return df
+        result = tool_query_product(product)
+
+        if not result:
+            return (
+                "No data found for this product.\n\n"
+                "Would you like me to create a support ticket?"
+            )
+        return result
 
     if "support" in text or "ticket" in text:
-        ticket = tool_create_ticket(user_input)
+        t = tool_create_support_ticket(user_input)
         return (
             f"🎫 **Support Ticket Created**\n"
-            f"ID: {ticket['ticket_id']}\n"
-            f"System: {ticket['system']}"
+            f"- ID: {t['ticket_id']}\n"
+            f"- System: {t['system']}\n"
+            f"- Status: {t['status']}"
         )
 
-    # LLM fallback (ONLY TEXT + SMALL CONTEXT)
+    # ---- LLM FALLBACK (NO RAW DATA PASSED) ----
     prompt = f"""
 You are an AI procurement assistant.
-You do NOT have direct DB access.
-If question seems unclear, suggest contacting support.
+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, temperature=0.3)
+    return client.text_generation(
+        prompt,
+        max_new_tokens=150,
+        temperature=0.3
+    )
 
-# ======================
+# ==================================================
 # STREAMLIT UI
-# ======================
+# ==================================================
 st.set_page_config(page_title="AI Procurement Agent", layout="wide")
-st.title("🤖 AI Procurement Agent")
+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.metric("Rows", stats["rows"])
 st.sidebar.metric("Avg Price", f"${stats['avg_price']}")
 st.sidebar.metric("Avg Delivery Days", stats["avg_delivery_days"])
@@ -125,9 +167,15 @@ database stats
 create support ticket
 """)
 
-# HF Token
-hf_token = st.sidebar.text_input("HF Token", type="password")
+# ---- HF Token ----
+hf_token = st.sidebar.text_input(
+    "Hugging Face Token",
+    type="password",
+    help="https://huggingface.co/settings/tokens"
+)
+
 if not hf_token:
+    st.warning("Please provide HF token")
     st.stop()
 
 client = InferenceClient(
@@ -135,11 +183,11 @@ client = InferenceClient(
     token=hf_token
 )
 
-# Chat
+# ---- Chat ----
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-user_input = st.chat_input("Ask the agent...")
+user_input = st.chat_input("Ask the procurement agent...")
 
 if user_input:
     st.session_state.messages.append(("user", user_input))