app.py

import json
import html
import pandas as pd
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"

ORDER_COLUMNS = [
    "customer",
    "item",
    "quantity",
    "flavor",
    "pickup_time",
    "delivery_address",
    "payment_status",
    "notes",
    "missing_fields",
]

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
model.eval()


SYSTEM_PROMPT = """
You are a careful order extraction engine for tiny sellers.

Extract customer orders from messy DMs. Return only valid JSON with this exact shape:
{
  "orders": [
    {
      "customer": "",
      "item": "",
      "quantity": "",
      "flavor": "",
      "pickup_time": "",
      "delivery_address": "",
      "payment_status": "",
      "notes": "",
      "missing_fields": []
    }
  ],
  "prep_list": [],
  "reply_drafts": []
}

Critical rules:
- Treat each line as one separate customer message.
- The text before the first ":" is the customer name.
- Copy customer names exactly as written. Do not uppercase or lowercase them.
- Never copy details from one customer's message into another customer's order.
- Include every customer message that looks like an order or possible order.
- Use only facts explicitly present in that customer's own message.
- If a value is unknown, use an empty string.
- Do not add order_id or total_cost.
- For pickup orders, put pickup time in pickup_time. Put a pickup place or delivery address in delivery_address.
- If the customer is unsure, still include the order and describe the uncertainty in notes.
- missing_fields should only include fields the seller needs to ask for: quantity, flavor, pickup_time, delivery_address, payment_status.
- Always set prep_list to [].
- Always set reply_drafts to [].
"""

SINGLE_ORDER_PROMPT = """
You extract one order from one customer's DM.

Return only valid JSON with this exact shape:
{
  "item": "",
  "quantity": "",
  "flavor": "",
  "pickup_time": "",
  "delivery_address": "",
  "payment_status": "",
  "notes": "",
  "missing_fields": []
}

Rules:
- Use only facts from this one message.
- Do not invent details.
- Put dates and times in pickup_time, such as "tomorrow", "Saturday morning", or "Friday 5pm".
- Put pickup places or delivery addresses in delivery_address, such as "farmers market".
- "pickup at the farmers market" means delivery_address is "farmers market", not pickup_time.
- "paid already" means payment_status is "paid".
- "I can pay Venmo" means payment_status is "can pay Venmo".
- If unknown, use an empty string.
- Do not ask for flavor unless the product clearly needs a flavor choice.
- missing_fields can only contain: quantity, flavor, pickup_time, delivery_address, payment_status.
"""


EXAMPLE_INPUT = """Maya: Hi! Can I get 2 dozen cupcakes for Saturday morning? Half vanilla, half chocolate.
Sam: Need 1 birthday cake, chocolate, for pickup Friday 5pm. I can pay Venmo.
Lena: Do you still have lemon bars? I need some for tomorrow but not sure how many yet.
Chris: 12 cookies please, pickup at the farmers market. Paid already.
"""

FOOD_TRUCK_EXAMPLE = """Alex: Can I get 3 chicken tacos for pickup at 6:30 tonight? Paid on Cash App.
Jamie: Do you still have vegan bowls? Need 2 tomorrow for office lunch.
Priya: One brisket sandwich, no onions. I'll pick up at the truck on Main Street.
Nate: 4 lemonades for the soccer team, pickup after practice.
"""

CRAFT_SELLER_EXAMPLE = """Olivia: I want 2 custom mugs with blue initials. Can you ship to 18 Pine Road?
Ben: Need one candle gift box for Saturday. Lavender if you have it.
Rosa: Can I order 3 tote bags? I can pick up at the market.
Eli: Do you still make birthday stickers? Need some next week but not sure how many.
"""

FARMERS_MARKET_EXAMPLE = """Grace: Can you hold 2 sourdough loaves for Sunday pickup?
Leo: I need 1 jar of strawberry jam and 2 honey bottles. Paid already.
Mina: Do you have eggs this weekend? Maybe 2 dozen if available.
Noah: Please save me 3 bags of granola, pickup at the farmers market.
"""

EXAMPLES = [
    [EXAMPLE_INPUT],
    [FOOD_TRUCK_EXAMPLE],
    [CRAFT_SELLER_EXAMPLE],
    [FARMERS_MARKET_EXAMPLE],
]

def extract_json(text):
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1:
        raise ValueError("No JSON object found")
    return json.loads(text[start:end + 1])

def normalize_orders(data):
    rows = []
    for order in data.get("orders", []):
        row = {}
        for col in ORDER_COLUMNS:
            value = order.get(col, "")
            if isinstance(value, list):
                value = ", ".join(str(v) for v in value)
            row[col] = value
        rows.append(row)
    return pd.DataFrame(rows, columns=ORDER_COLUMNS)

def format_list(title, items):
    if not items:
        return f"### {title}\nNothing found."
    lines = []
    for item in items:
        if isinstance(item, dict):
            lines.append("- " + json.dumps(item, ensure_ascii=False))
        else:
            lines.append(f"- {item}")
    return f"### {title}\n" + "\n".join(lines)

def format_replies(replies):
    if not replies:
        return "### Reply drafts\nNothing found."
    lines = []
    for reply in replies:
        customer = reply.get("customer", "Customer")
        text = reply.get("reply", "")
        lines.append(f"**{customer}**\n\n{text}")
    return "### Reply drafts\n\n" + "\n\n---\n\n".join(lines)

def text_value(value):
    if isinstance(value, list):
        return ", ".join(str(v) for v in value if str(v).strip())
    if value is None:
        return ""
    return str(value).strip()

def missing_list(order):
    raw = order.get("missing_fields", [])
    if isinstance(raw, str):
        fields = [part.strip() for part in raw.split(",") if part.strip()]
    else:
        fields = [str(part).strip() for part in raw if str(part).strip()]

    allowed = {"quantity", "flavor", "pickup_time", "delivery_address", "payment_status"}
    fields = [field for field in fields if field in allowed]

    item = text_value(order.get("item"))
    quantity = text_value(order.get("quantity"))
    flavor = text_value(order.get("flavor"))
    pickup_time = text_value(order.get("pickup_time"))
    delivery_address = text_value(order.get("delivery_address"))
    payment_status = text_value(order.get("payment_status"))

    if item and not quantity:
        fields.append("quantity")

    if pickup_time:
        fields = [field for field in fields if field != "pickup_time"]
    if delivery_address:
        fields = [field for field in fields if field != "delivery_address"]
    if payment_status:
        fields = [field for field in fields if field != "payment_status"]
    if flavor:
        fields = [field for field in fields if field != "flavor"]

    if "flavor" in fields and item.lower() not in ["cake", "birthday cake", "cupcakes"]:
        fields = [field for field in fields if field != "flavor"]

    return sorted(set(fields))

def post_process_order(order, message):
    msg = message.lower()

    if "paid already" in msg or "already paid" in msg:
        order["payment_status"] = "paid"
    elif "venmo" in msg:
        order["payment_status"] = "can pay Venmo"
    elif "paid" not in msg and "venmo" not in msg:
        order["payment_status"] = ""

    pickup_time = text_value(order.get("pickup_time"))
    if "paid" in pickup_time.lower() or "venmo" in pickup_time.lower():
        order["pickup_time"] = ""

    if "farmers market" in msg:
        order["delivery_address"] = "farmers market"
        if "farmers market" in text_value(order.get("pickup_time")).lower():
            order["pickup_time"] = ""

    order["missing_fields"] = missing_list(order)
    return order

def build_prep_list(data):
    items = []
    for order in data.get("orders", []):
        item = text_value(order.get("item"))
        if not item:
            continue

        customer = text_value(order.get("customer")) or "customer"
        quantity = text_value(order.get("quantity")) or "quantity to confirm"
        flavor = text_value(order.get("flavor"))

        line = f"{quantity} {item}"
        if flavor:
            line += f" ({flavor})"
        line += f" - {customer}"
        items.append(line)

    return items

def build_reply_drafts(data):
    replies = []
    labels = {
        "quantity": "quantity",
        "flavor": "flavor",
        "pickup_time": "pickup or delivery time",
        "delivery_address": "pickup place",
        "payment_status": "payment status",
    }

    for order in data.get("orders", []):
        customer = text_value(order.get("customer")) or "there"
        item = text_value(order.get("item")) or "order"
        quantity = text_value(order.get("quantity"))
        flavor = text_value(order.get("flavor"))
        missing = [labels.get(field, field) for field in missing_list(order)]

        if missing:
            needed = ", ".join(missing)
            reply = f"Thanks, {customer}! I have your {item} order. Could you confirm the {needed}?"
        else:
            summary = f"{quantity} {item}".strip()
            if flavor:
                summary += f" ({flavor})"
            reply = f"Thanks, {customer}! Confirming your order: {summary}."

        replies.append({"customer": customer, "reply": reply})

    return replies

def build_summary(data):
    orders = data.get("orders", [])
    total = len(orders)
    followups = sum(1 for order in orders if missing_list(order))
    ready = total - followups

    return f"""
            <div style="display:grid;grid-template-columns:repeat(3,minmax(0,1fr));gap:12px;margin:8px 0 18px;">
              <div style="background:#f8fafc;border:1px solid #e2e8f0;border-radius:10px;padding:14px;">
                <div style="font-size:24px;font-weight:700;">{total}</div>
                <div style="color:#475569;">customer messages</div>
              </div>
              <div style="background:#ecfdf5;border:1px solid #bbf7d0;border-radius:10px;padding:14px;">
                <div style="font-size:24px;font-weight:700;">{ready}</div>
                <div style="color:#166534;">ready orders</div>
              </div>
              <div style="background:#fffbeb;border:1px solid #fde68a;border-radius:10px;padding:14px;">
                <div style="font-size:24px;font-weight:700;">{followups}</div>
                <div style="color:#92400e;">follow-ups needed</div>
              </div>
            </div>
            """

def build_review_board(data):
    ready_cards = []
    followup_cards = []

    for order in data.get("orders", []):
        customer = html.escape(text_value(order.get("customer")) or "Customer")
        item = html.escape(text_value(order.get("item")) or "order")
        quantity = html.escape(text_value(order.get("quantity")) or "quantity to confirm")
        flavor = html.escape(text_value(order.get("flavor")))
        pickup_time = html.escape(text_value(order.get("pickup_time")))
        delivery_address = html.escape(text_value(order.get("delivery_address")))
        payment_status = html.escape(text_value(order.get("payment_status")))
        fields = missing_list(order)

        title = f"{quantity} {item}"
        if flavor:
            title += f" ({flavor})"

        details = []
        if pickup_time:
            details.append(f"Time: {pickup_time}")
        if delivery_address:
            details.append(f"Place: {delivery_address}")
        if payment_status:
            details.append(f"Payment: {payment_status}")

        detail_html = " · ".join(details) if details else "Details still need review."

        if fields:
            chips = "".join(
                f'<span style="display:inline-block;background:#fef3c7;color:#92400e;border-radius:999px;padding:4px 9px;margin:3px;font-size:12px;">{html.escape(field.replace("_", " "))}</span>'
                for field in fields
            )
            needed = ", ".join(field.replace("_", " ") for field in fields)
            suggested_reply = html.escape(
                f"Thanks, {customer}! I have your {item} order. Could you confirm the {needed}?"
            )
            
            followup_cards.append(
                f"""
            <div style="background:#fff7ed;border:1px solid #fed7aa;border-radius:10px;padding:14px;margin-bottom:10px;">
              <div style="font-weight:700;">{customer}</div>
              <div>{title}</div>
              <div style="color:#64748b;font-size:13px;margin-top:4px;">{detail_html}</div>
              <div style="margin-top:8px;">{chips}</div>
              <div style="margin-top:10px;background:#ffffff;border:1px solid #fed7aa;border-radius:8px;padding:10px;color:#334155;font-size:13px;">
                <div style="font-weight:700;margin-bottom:4px;">Suggested reply</div>
                {suggested_reply}
              </div>
            </div>
            """
            )
        else:
            ready_cards.append(
                f"""
                <div style="background:#f0fdf4;border:1px solid #bbf7d0;border-radius:10px;padding:14px;margin-bottom:10px;">
                  <div style="font-weight:700;">{customer}</div>
                  <div>{title}</div>
                  <div style="color:#64748b;font-size:13px;margin-top:4px;">{detail_html}</div>
                </div>
                """
            )

    ready_html = "".join(ready_cards) or '<div style="color:#64748b;">No ready orders yet.</div>'
    followup_html = "".join(followup_cards) or '<div style="color:#64748b;">No follow-ups needed.</div>'

    return f"""
            <div style="display:grid;grid-template-columns:1fr 1fr;gap:16px;">
              <section>
                <h3>Ready to prep</h3>
                {ready_html}
              </section>
              <section>
                <h3>Needs follow-up</h3>
                {followup_html}
              </section>
            </div>
            """

def split_customer_messages(messages):
    entries = []
    current_customer = ""
    current_parts = []

    for raw_line in messages.splitlines():
        line = raw_line.strip()
        if not line:
            continue

        if ":" in line:
            possible_name, body = line.split(":", 1)
            if possible_name.strip() and len(possible_name.strip().split()) <= 3:
                if current_customer or current_parts:
                    entries.append((current_customer or "Customer", " ".join(current_parts).strip()))
                current_customer = possible_name.strip()
                current_parts = [body.strip()]
                continue

        if current_parts:
            current_parts.append(line)
        else:
            entries.append(("Customer", line))

    if current_customer or current_parts:
        entries.append((current_customer or "Customer", " ".join(current_parts).strip()))

    return [(name, body) for name, body in entries if body]

def extract_single_order(customer, message):
    prompt = tokenizer.apply_chat_template(
        [
            {"role": "system", "content": SINGLE_ORDER_PROMPT},
            {"role": "user", "content": f"Customer: {customer}\nMessage: {message}"},
        ],
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=350,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated = tokenizer.decode(
        output[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True,
    )

    try:
        parsed = extract_json(generated)
    except Exception:
        parsed = {
            "item": "",
            "quantity": "",
            "flavor": "",
            "pickup_time": "",
            "delivery_address": "",
            "payment_status": "",
            "notes": message,
            "missing_fields": [],
        }

    order = {"customer": customer}
    for col in ORDER_COLUMNS[1:]:
        value = parsed.get(col, "")
        if col == "missing_fields":
            if isinstance(value, list):
                order[col] = value
            elif isinstance(value, str):
                order[col] = [part.strip() for part in value.split(",") if part.strip()]
            else:
                order[col] = []
        else:
            order[col] = text_value(value)

    return post_process_order(order, message)

def analyze_messages(messages):
    if not messages.strip():
        return pd.DataFrame(columns=ORDER_COLUMNS), "Choose a scenario and organize the inbox.", "", "", "", ""

    entries = split_customer_messages(messages)
    orders_data = [extract_single_order(customer, message) for customer, message in entries]

    data = {"orders": orders_data}
    orders_df = normalize_orders(data)

    auto_prep = build_prep_list(data)
    auto_replies = build_reply_drafts(data)

    data["prep_list"] = auto_prep
    data["reply_drafts"] = auto_replies

    summary = build_summary(data)
    review = build_review_board(data)
    prep = format_list("Prep list", auto_prep)
    replies = format_replies(auto_replies)
    raw = json.dumps(data, indent=2, ensure_ascii=False)
    return orders_df, summary, review, prep, replies, raw

with gr.Blocks(title="DM Order Desk", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        f"""
        # DM Order Desk
        
        A small-model inbox assistant for tiny sellers.
        
        **Unofficial Build Small Hackathon-inspired project** · `{MODEL_ID}` · Gradio Space
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Customer inbox")
            messages = gr.Textbox(
                label="Messy customer DMs",
                value=EXAMPLE_INPUT,
                lines=14,
            )
            run = gr.Button("Organize orders", variant="primary")

            gr.Markdown("### Demo scenarios")
            with gr.Row():
                home_btn = gr.Button("Home Bakery")
                food_btn = gr.Button("Food Truck")
            with gr.Row():
                craft_btn = gr.Button("Craft Market")
                farm_btn = gr.Button("Farmers Market")

        with gr.Column(scale=2):
            summary = gr.HTML("Choose a scenario and organize the inbox.")

            with gr.Tabs():
                with gr.Tab("Review Board"):
                    review = gr.HTML()
                with gr.Tab("Data Table"):
                    orders = gr.Dataframe(label="Orders", headers=ORDER_COLUMNS)
                with gr.Tab("Prep List"):
                    prep = gr.Markdown()
                with gr.Tab("Reply Drafts"):
                    replies = gr.Markdown()
                with gr.Tab("Raw JSON"):
                    raw = gr.Code(label="Raw JSON", language="json")

    run.click(analyze_messages, inputs=messages, outputs=[orders, summary, review, prep, replies, raw])

    home_btn.click(lambda: EXAMPLE_INPUT, outputs=messages)
    food_btn.click(lambda: FOOD_TRUCK_EXAMPLE, outputs=messages)
    craft_btn.click(lambda: CRAFT_SELLER_EXAMPLE, outputs=messages)
    farm_btn.click(lambda: FARMERS_MARKET_EXAMPLE, outputs=messages)

demo.launch()