app/tasks/order_comms.py

"""Composite support-comms task: one customer communication, two LLM jobs.

Given one inbound communication (email / chat / call), the model must do BOTH:
  1. classify the inquiry intent into a two-level taxonomy (category -> subcategory), and
  2. tag the business entities it references (Order/Shipment/Return + Customer/Product),
     using two tools over a committed synthetic world: a fuzzy entity `search_entities`
     index and an executable `execute_sql` (duckdb) over the `shop.*` schema.

Ground truth is by construction: a seed-picked SCENARIO jointly fixes the gold intent,
the gold entity set, and a consistent rendered communication. This re-skins a real
field-service customer workload onto a fully generic e-commerce surface — same shape,
no associable detail. See docs/superpowers/specs/2026-06-18-order-comms-design.md.
"""

import json
import random
import re

import duckdb

from app.config import DATA_DIR
from app.tasks.base import Sample, Task

WORLD_DIR = DATA_DIR / "order_world"

# --- Taxonomy (Task 1) -------------------------------------------------------

TAXONOMY = {
    "Shipping": ["late_delivery", "lost_package", "wrong_address", "tracking_request", "general"],
    "Returns": ["damaged_item", "wrong_item", "changed_mind", "return_status", "general"],
    "Billing": ["double_charge", "refund_status", "coupon_promo", "payment_failed", "general"],
    "Account": ["login_issue", "update_info", "close_account", "general"],
    "Product": ["availability", "how_to_use", "specs_question", "general"],
    "Order_Change": ["cancel_order", "modify_order", "address_change", "general"],
}

# --- Load the committed world ------------------------------------------------


def _load(name: str) -> list[dict]:
    with (WORLD_DIR / f"{name}.jsonl").open() as f:
        return [json.loads(line) for line in f if line.strip()]


CUSTOMERS = _load("customers")
PRODUCTS = _load("products")
ORDERS = _load("orders")
ORDER_ITEMS = _load("order_items")
SHIPMENTS = _load("shipments")
RETURNS = _load("returns")

CUST_BY_ID = {c["customer_id"]: c for c in CUSTOMERS}
PROD_BY_SKU = {p["sku"]: p for p in PRODUCTS}
ITEMS_BY_ORDER: dict[str, list[dict]] = {}
for _it in ORDER_ITEMS:
    ITEMS_BY_ORDER.setdefault(_it["order_id"], []).append(_it)
SHIP_BY_ORDER: dict[str, list[dict]] = {}
for _s in SHIPMENTS:
    SHIP_BY_ORDER.setdefault(_s["order_id"], []).append(_s)
RET_BY_ORDER: dict[str, list[dict]] = {}
for _r in RETURNS:
    RET_BY_ORDER.setdefault(_r["order_id"], []).append(_r)
ORDERS_BY_CUST: dict[str, list[dict]] = {}
for _o in ORDERS:
    ORDERS_BY_CUST.setdefault(_o["customer_id"], []).append(_o)

# Every valid id in the world — used by score() for the hallucination check.
WORLD_FRIENDLY_IDS = (
    {o["order_number"].upper() for o in ORDERS}
    | {s["shipment_number"].upper() for s in SHIPMENTS}
    | {r["rma_number"].upper() for r in RETURNS}
    | {p["sku"].upper() for p in PRODUCTS}
)
WORLD_UUIDS = (
    {o["order_id"] for o in ORDERS}
    | {s["shipment_id"] for s in SHIPMENTS}
    | {r["return_id"] for r in RETURNS}
    | {c["customer_id"] for c in CUSTOMERS}
)


def _product_phrase(prod: dict) -> str:
    """The product name with the brand stripped, e.g. 'espresso machine'."""
    name, brand = prod["product_name"], prod["brand"]
    base = name[len(brand):].strip() if name.startswith(brand) else name
    return base.lower()


def _main_product(order_id: str) -> dict:
    items = ITEMS_BY_ORDER.get(order_id, [])
    return PROD_BY_SKU[items[0]["sku"]] if items else PRODUCTS[0]


# --- Tool 1: search_entities (mock fuzzy index) ------------------------------


def _tok(s: str) -> list[str]:
    return re.findall(r"[a-z0-9]+", (s or "").lower())


def _build_index() -> list[dict]:
    docs: list[dict] = []

    def add(eid: str, etype: str, exact: str, text_parts: list[str], fields: dict) -> None:
        text = " ".join(str(p) for p in text_parts if p)
        docs.append({
            "id": eid, "type": etype, "exact": (exact or "").upper(),
            "tokens": set(_tok(text)), "text": text.lower(), "fields": fields,
        })

    for o in ORDERS:
        cust = CUST_BY_ID[o["customer_id"]]
        prods = [PROD_BY_SKU[i["sku"]] for i in ITEMS_BY_ORDER.get(o["order_id"], [])]
        pnames = [p["product_name"] for p in prods]
        add(o["order_id"], "Order", o["order_number"],
            [o["order_number"], cust["full_name"], o["ship_to_city"], o["ship_to_state"], *pnames],
            {"order_number": o["order_number"], "status": o["order_status"],
             "customer_name": cust["full_name"],
             "ship_to": f"{o['ship_to_city']}, {o['ship_to_state']}", "total": o["order_total"]})
    for s in SHIPMENTS:
        onum = next((o["order_number"] for o in ORDERS if o["order_id"] == s["order_id"]), "")
        add(s["shipment_id"], "Shipment", s["shipment_number"],
            [s["shipment_number"], s["carrier"], s["tracking_number"], onum],
            {"shipment_number": s["shipment_number"], "carrier": s["carrier"],
             "status": s["ship_status"], "tracking_number": s["tracking_number"],
             "order_number": onum})
    for r in RETURNS:
        onum = next((o["order_number"] for o in ORDERS if o["order_id"] == r["order_id"]), "")
        add(r["return_id"], "Return", r["rma_number"],
            [r["rma_number"], onum, r["reason"], r["sku"]],
            {"rma_number": r["rma_number"], "status": r["return_status"],
             "reason": r["reason"], "order_number": onum})
    for p in PRODUCTS:
        add(p["sku"], "Product", p["sku"],
            [p["product_name"], p["brand"], p["category"], p["sku"]],
            {"sku": p["sku"], "product_name": p["product_name"],
             "brand": p["brand"], "category": p["category"]})
    for c in CUSTOMERS:  # NOTE: phone is deliberately NOT indexed (forces SQL)
        add(c["customer_id"], "Customer", "",
            [c["full_name"], c["email"], c["city"]],
            {"name": c["full_name"], "email": c["email"], "city": c["city"]})
    return docs


INDEX = _build_index()


def search_entities(search_terms: str, entity_types: list[str] | None = None,
                    operator: str = "And", max_results: int = 10) -> str:
    terms = (search_terms or "").strip()
    if len(terms) < 3:
        return json.dumps({"error": "search error: INVALID_ARGUMENT — "
                           "search term must contain at least 3 characters"})
    qtokens = [t for t in _tok(terms) if len(t) >= 3]
    want = {t for t in (entity_types or [])}
    by_type: dict[str, list[tuple[float, dict]]] = {}
    upper = terms.upper()
    for doc in INDEX:
        if want and doc["type"] not in want:
            continue
        score = 100.0 if doc["exact"] and doc["exact"] in upper else 0.0
        for t in qtokens:
            if t in doc["tokens"]:
                score += 3.0
            elif t in doc["text"]:
                score += 1.0
        if score > 0:
            by_type.setdefault(doc["type"], []).append((score, doc))

    cap = max(1, min(int(max_results or 10), 50))
    results = []
    for lst in by_type.values():
        lst.sort(key=lambda sd: (-sd[0], sd[1]["id"]))
        for score, doc in lst[:cap]:
            results.append({"id": doc["id"], "type": doc["type"],
                            "score": round(score, 4), "fields": doc["fields"]})
    if not results:
        return json.dumps({"error": "search error: NOT_FOUND — No results found"})
    results.sort(key=lambda r: (-r["score"], r["id"]))
    return json.dumps({"result_count": len(results), "results": results})


# --- Tool 2: execute_sql (duckdb over shop.*) --------------------------------

_DB = None
_FORBIDDEN = re.compile(
    r"\b(insert|update|delete|drop|create|alter|attach|copy|pragma|install|load|set)\b", re.I)


def _db() -> duckdb.DuckDBPyConnection:
    global _DB
    if _DB is None:
        con = duckdb.connect(":memory:")
        con.execute("CREATE SCHEMA shop")
        for name in ("customers", "products", "orders", "order_items", "shipments", "returns"):
            path = str(WORLD_DIR / f"{name}.jsonl")
            con.execute(f"CREATE TABLE shop.{name} AS "
                        f"SELECT * FROM read_json_auto('{path}', format='newline_delimited')")
        con.execute("SET search_path='shop'")  # bare table names also resolve
        _DB = con
    return _DB


def execute_sql(sql: str, reason: str = "") -> str:
    query = (sql or "").strip().rstrip(";").strip()
    # normalize a warehouse-style prefix if the model defaults to one
    query = re.sub(r"(?i)\b(prod\.)?edw\.", "shop.", query)
    if ";" in query:
        return json.dumps({"error": "SQL error: only a single statement is allowed"})
    if not re.match(r"(?i)^\s*(select|with)\b", query):
        return json.dumps({"error": "SQL error: only SELECT/WITH queries are allowed"})
    if _FORBIDDEN.search(query):
        return json.dumps({"error": "SQL error: only read-only SELECT/WITH queries are allowed"})
    if not re.search(r"(?i)\blimit\b", query):
        query += " LIMIT 50"
    try:
        cur = _db().execute(query)
        columns = [d[0] for d in cur.description]
        rows = [dict(zip(columns, r)) for r in cur.fetchmany(50)]
    except Exception as e:  # noqa: BLE001 — surface the DB error to the agent, never raise
        return json.dumps({"error": f"SQL error: {str(e).splitlines()[0][:300]}"})
    return json.dumps({"row_count": len(rows), "columns": columns, "rows": rows}, default=str)


def execute_tool(name: str, args: dict) -> str:
    if name == "search_entities":
        return search_entities(
            args.get("search_terms", ""), args.get("entity_types"),
            args.get("operator", "And"), args.get("max_results", 10))
    if name == "execute_sql":
        return execute_sql(args.get("sql", ""), args.get("reason", ""))
    return json.dumps({"error": f"unknown tool {name}"})


TOOLS = [
    {"type": "function", "function": {
        "name": "search_entities",
        "description": (
            "Fuzzy text search across shop entities (Order, Shipment, Return, Product, "
            "Customer). WORKS WELL FOR: full friendly IDs (ORD-/SHP-/RMA-, exact-match "
            "boosted), product/brand names, customer names, emails, cities, tracking "
            "numbers. DOES NOT WORK FOR: phone numbers (not indexed) or UUIDs — use "
            "execute_sql for those. Returns a JSON string."),
        "parameters": {"type": "object", "properties": {
            "search_terms": {"type": "string", "description": "Distinctive words to search for. Full IDs like 'ORD-260512-0473' work best."},
            "reason": {"type": "string", "description": "Brief explanation of what you're looking for."},
            "entity_types": {"type": "array", "items": {"type": "string"},
                             "description": "Subset of Order, Shipment, Return, Product, Customer. Omit to search all."},
            "operator": {"type": "string", "default": "And", "description": "'And' or 'Or'."},
            "max_results": {"type": "integer", "default": 10, "description": "Max results per type (max 50)."},
        }, "required": ["search_terms", "reason"]}}},
    {"type": "function", "function": {
        "name": "execute_sql",
        "description": (
            "Run a read-only SQL query (SELECT/WITH only) against the shop warehouse "
            "(duckdb). Use for phone/email lookups, UUID joins, date filtering, and "
            "anything search can't do. Returns up to 50 rows as a JSON string."),
        "parameters": {"type": "object", "properties": {
            "sql": {"type": "string", "description": "A single SELECT statement. Use shop.<table>."},
            "reason": {"type": "string", "description": "Brief explanation of what you're looking for."},
        }, "required": ["sql", "reason"]}}},
]

# --- System prompt -----------------------------------------------------------

SYSTEM_PROMPT = """\
You are a support-operations analyst for a large online retail marketplace. You are \
given ONE inbound customer communication (email, chat, or call transcript). Do TWO \
things and return a single JSON object.

## Task 1 — Classify the inquiry intent
Pick exactly one (category, subcategory) from this taxonomy. Use the category's \
"general" subcategory only when the message clearly belongs to the category but the \
specific subcategory cannot be pinpointed.
- Shipping: late_delivery, lost_package, wrong_address, tracking_request, general
- Returns: damaged_item, wrong_item, changed_mind, return_status, general
- Billing: double_charge, refund_status, coupon_promo, payment_failed, general
- Account: login_issue, update_info, close_account, general
- Product: availability, how_to_use, specs_question, general
- Order_Change: cancel_order, modify_order, address_change, general

## Task 2 — Tag the business entities the communication is about
Find the transaction(s) this message concerns — Order (ORD-…), Shipment (SHP-…), or \
Return (RMA-…) — plus any supporting Customer or Product. Every work-related message \
has at least one transaction; returning none when one exists is a miss. For each tagged \
entity provide its `entity_id` (the UUID) — you can only get the UUID from a tool result, \
so look it up; do not invent it.

## Tools
1. search_entities — fuzzy text search. Best for full friendly IDs (exact-match boosted), \
product/brand names, customer names, emails, cities. NOT for phone numbers (not indexed) \
or UUIDs — use SQL for those.
2. execute_sql — read-only SQL (SELECT/WITH only) over the `shop` schema (duckdb). \
Use for phone/email lookups, UUID joins, and date filtering.

### shop schema (do not guess columns; this is the whole schema)
- shop.customers(customer_id, full_name, email, phone, city, state_code)
- shop.products(sku, product_name, brand, category, price)
- shop.orders(order_id, order_number, customer_id, order_status, order_total, placed_at, ship_to_city, ship_to_state)
- shop.order_items(order_id, sku, quantity, unit_price)
- shop.shipments(shipment_id, shipment_number, order_id, carrier, ship_status, tracking_number, shipped_at, est_delivery, delivered_at)
- shop.returns(return_id, rma_number, order_id, sku, reason, return_status, created_at)
phone is E.164 ('+1XXXXXXXXXX') in shop.customers and is NOT searchable — match it in SQL.

### Canonical lookups
- ID in the message → search_entities for the full ORD-/SHP-/RMA- number.
- Call with no ID → resolve the caller:
    SELECT customer_id, full_name FROM shop.customers WHERE phone = '<+1XXXXXXXXXX>'
  then that customer's recent orders (with main product + shipment):
    SELECT o.order_number, o.order_id, o.order_status, o.placed_at, p.product_name, s.shipment_number, s.ship_status
    FROM shop.orders o
    JOIN shop.order_items oi ON oi.order_id = o.order_id
    JOIN shop.products p ON p.sku = oi.sku
    LEFT JOIN shop.shipments s ON s.order_id = o.order_id
    WHERE o.customer_id = '<uuid>' ORDER BY o.placed_at DESC

## Pre-Fetched Context / Existing References
If the message has a "### Pre-Fetched Context" section, those candidates were resolved for \
you (an ID in the body, or the sender's recent orders, or an email-thread reference). It is \
a CANDIDATE LIST, not the answer — for a phone-resolved list, tag only the order(s) this \
conversation is actually about (match on product, issue, and timing); do NOT tag the whole \
list. An "### Existing References" section lists entities upstream systems already attached \
— "(Explicit)" are near-confirmed; verify rather than re-discover.

## Budget
Aim to finish within a few tool calls; you may call tools in parallel. Produce your answer \
as soon as you have enough evidence.

## Output (exactly one fenced ```json block, nothing after it)
```json
{
  "classification": {"category": "<Category>", "subcategory": "<subcategory>"},
  "entities": [
    {"entity_type": "Order|Shipment|Return|Product|Customer", "friendly_id": "ORD-…",
     "entity_id": "<uuid>", "confidence": 0.0, "reasoning": "…"}
  ],
  "search_summary": "what you searched and why"
}
```\
"""

# --- Scenario-driven episode generation --------------------------------------

# Pre-bucket the world so a scenario can pick a target that actually fits.
_ORD_W_SHIP = [o for o in ORDERS if SHIP_BY_ORDER.get(o["order_id"])]
_ORD_W_RETURN = [o for o in ORDERS if RET_BY_ORDER.get(o["order_id"])]
_ORD_INTRANSIT = [o for o in _ORD_W_SHIP
                  if SHIP_BY_ORDER[o["order_id"]][0]["ship_status"] in
                  ("in_transit", "out_for_delivery", "exception", "label_created")]
_ORD_ACTIVE = [o for o in ORDERS if o["order_status"] in ("placed", "paid", "shipped")]

# Each scenario: (category, subcategory, primary_kind, pool, [transcript templates], [body templates])
# {p}=product phrase, {n}=days ago, {ord}=order number, {city}=ship-to city
SCENARIOS = [
    ("Shipping", "late_delivery", "shipment", _ORD_INTRANSIT,
     ["Hi, I ordered the {p} about {n} days ago and the tracking hasn't moved in days. "
      "Any idea when it'll actually get here?",
      "Yeah, I'm calling about a {p} I'm still waiting on — it was supposed to be here by "
      "now and the status hasn't changed. Can you check on it?"],
     ["My {p} order is way past the delivery estimate and tracking is stuck. Where is it?",
      "It's been {n} days and the {p} still hasn't shipped out properly — what's going on?"]),
    ("Shipping", "lost_package", "shipment", _ORD_INTRANSIT,
     ["The tracking says my {p} was delivered but it's nowhere — I think the package is lost. "
      "Can you help?",
      "I never got the {p} I ordered. Tracking went dead halfway. Pretty sure it's lost."],
     ["Package with my {p} never arrived even though it shows in transit forever. I think it's lost.",
      "My {p} shipment seems lost — no movement for over a week and nothing at my door."]),
    ("Shipping", "wrong_address", "order", _ORD_W_SHIP,
     ["I think my {p} is being sent to my old address. Can you change where it's going?",
      "The {p} order — I realized the shipping address is wrong. It's headed to the wrong place."],
     ["Please fix the shipping address on my {p} order, it's going to the wrong house.",
      "The address on my recent {p} order is outdated — can you reroute it?"]),
    ("Shipping", "tracking_request", "shipment", _ORD_W_SHIP,
     ["Can you just give me the tracking number for my {p}? I can't find it.",
      "I'm calling to get a tracking update on the {p} I ordered."],
     ["Can you send me the tracking link for my {p} order?",
      "What's the tracking status on my {p}? Haven't gotten an update."]),
    ("Returns", "damaged_item", "return", _ORD_W_RETURN,
     ["My {p} showed up cracked and unusable. I already started a return — what are next steps?",
      "Calling about a {p} that arrived damaged. I opened a return, want to make sure it's moving."],
     ["The {p} arrived broken. I filed a return — can you check the status?",
      "My {p} came damaged in the box. There's a return open on it, what now?"]),
    ("Returns", "wrong_item", "order", _ORD_W_SHIP,
     ["I ordered a {p} but the box had something completely different inside.",
      "The {p} I ordered? You shipped me the wrong thing entirely."],
     ["Wrong item — I ordered a {p} and received something else. Need to sort this out.",
      "You sent the wrong product for my {p} order."]),
    ("Returns", "return_status", "return", _ORD_W_RETURN,
     ["I sent back the {p} a while ago and haven't heard anything. Where's my return at?",
      "Checking on a return I opened for a {p} — what's the current status?"],
     ["What's the status of my return for the {p}? Shipped it back already.",
      "Any update on the {p} return I filed?"]),
    ("Billing", "double_charge", "order", ORDERS,
     ["I think I got charged twice for my {p} order. Can you look into the billing?",
      "There are two charges on my card for the same {p} order. Need a refund for one."],
     ["My card shows two charges for the {p} order — please refund the duplicate.",
      "Double charged on the {p} order. Can you fix the billing?"]),
    ("Billing", "refund_status", "return", _ORD_W_RETURN,
     ["I returned the {p} and was told I'd get a refund. Hasn't shown up yet — where is it?",
      "Calling about a refund for a {p} I sent back. When will I see the money?"],
     ["Still waiting on my refund for the returned {p}. Any timeline?",
      "Where's my refund for the {p} return? It's been a while."]),
    ("Order_Change", "cancel_order", "order", _ORD_ACTIVE,
     ["I need to cancel my {p} order — changed my mind before it ships.",
      "Please cancel the {p} I just ordered, I don't need it anymore."],
     ["Cancel my {p} order please, I don't want it after all.",
      "Need to cancel the recent {p} order before it goes out."]),
    ("Order_Change", "modify_order", "order", _ORD_ACTIVE,
     ["Can I change the quantity on my {p} order? I want to add one more.",
      "I'd like to modify my {p} order — swap one of the items if it hasn't shipped."],
     ["Please update my {p} order — I need to change what's in it.",
      "Can you modify the {p} order I placed? Want to tweak it."]),
    ("Product", "availability", "product", ORDERS,
     ["Quick question — is the {p} back in stock? Looking to buy another.",
      "Do you know if the {p} is available again? I had one and want a second."],
     ["Is the {p} in stock right now? Want to reorder.",
      "When will the {p} be available again?"]),
    # Ambiguous -> gold subcategory "general"
    ("Shipping", "general", "order", _ORD_W_SHIP,
     ["Hey, I'm calling about my {p} order — something's just not right with it, can you take a look?",
      "I have a question about the {p} I ordered, it's a bit of a mess, hoping you can help."],
     ["Something's off with my {p} order — not sure what exactly. Can you check the whole thing?",
      "Need help with my {p} order in general, a few things seem wrong."]),
]
_WEIGHTS = [3, 2, 2, 2, 3, 2, 2, 3, 2, 3, 2, 2, 2]  # len == len(SCENARIOS)


def _ent(etype: str, friendly_id: str | None, entity_id: str) -> dict:
    return {"entity_type": etype, "friendly_id": friendly_id, "entity_id": entity_id}


def generate(seed: int) -> tuple[dict, str]:
    """Deterministically build (ground-truth dict, rendered communication markdown)."""
    rng = random.Random(seed)
    category, sub, primary_kind, pool, calls, bodies = rng.choices(SCENARIOS, _WEIGHTS)[0]
    order = rng.choice(pool or ORDERS)
    oid = order["order_id"]
    cust = CUST_BY_ID[order["customer_id"]]
    prod = _main_product(oid)
    pphrase = _product_phrase(prod)
    ship = (SHIP_BY_ORDER.get(oid) or [None])[0]
    ret = (RET_BY_ORDER.get(oid) or [None])[0]

    # Gold: primary transaction + required + optional supporting entities
    order_ent = _ent("Order", order["order_number"], oid)
    if primary_kind == "shipment" and ship:
        primary = _ent("Shipment", ship["shipment_number"], ship["shipment_id"])
        required = [primary, order_ent]
    elif primary_kind == "return" and ret:
        primary = _ent("Return", ret["rma_number"], ret["return_id"])
        required = [primary, order_ent]
    elif primary_kind == "product":
        primary = _ent("Product", prod["sku"], prod["sku"])
        required = [primary]
    else:  # order (incl. shipment/return scenarios whose child happens to be missing)
        primary = order_ent
        required = [order_ent]
    optional = [_ent("Customer", None, cust["customer_id"]),
                _ent("Product", prod["sku"], prod["sku"])]
    if ship and primary["entity_id"] != ship["shipment_id"]:
        optional.append(_ent("Shipment", ship["shipment_number"], ship["shipment_id"]))

    # Lead type drives channel + difficulty
    lead = rng.choices(["explicit", "thread", "phone"], [35, 25, 40])[0]
    days = rng.randint(4, 14)

    def order_line(o: dict, indent: str = "    ") -> str:
        p = _product_phrase(_main_product(o["order_id"]))
        s = (SHIP_BY_ORDER.get(o["order_id"]) or [None])[0]
        tail = f", Shipment {s['shipment_number']} {s['ship_status']}" if s else ""
        return (f"{indent}Order {o['order_number']} — {p}, placed {o['placed_at'][:10]}, "
                f"status {o['order_status']}{tail}")

    prefetch, existing = [], []
    if lead == "explicit":
        channel = rng.choice(["Email", "Chat"])
        shown_id = primary["friendly_id"] or order["order_number"]
        body = rng.choice(bodies).format(p=pphrase, n=days)
        body += f" Order number is {order['order_number']}."
        if primary["friendly_id"] and not primary["friendly_id"].startswith("ORD"):
            body += f" Reference {primary['friendly_id']}."
        prefetch = ["These transactions were pre-resolved from IDs in the message. "
                    "Verify against the message before tagging.",
                    f"- {shown_id} (in message) → {order_line(order, '').strip()}"]
        content = body
    elif lead == "thread":
        channel = "Email"
        content = (f"Subject: Re: your order {order['order_number']}\n\n"
                   + rng.choice(bodies).format(p=pphrase, n=days))
        prefetch = ["These were pre-resolved from this email thread.",
                    f"- Earlier in this thread: Order {order['order_number']} — a reply usually "
                    "concerns the same order; confirm against this message and tag if it applies."]
        existing = [f"- Order: {order['order_number']} (from Friendly ID Matcher (Explicit))"]
    else:  # phone
        channel = "Call"
        content = ("Agent: Thanks for calling, how can I help?\n"
                   f"Caller: {rng.choice(calls).format(p=pphrase, n=days)}\n"
                   "Agent: Sure, let me pull up your account and take a look.")
        # candidate list: target + a few of the customer's other orders as distractors
        others = [o for o in ORDERS_BY_CUST.get(cust["customer_id"], []) if o["order_id"] != oid]
        rng.shuffle(others)
        candidates = sorted([order, *others[:3]], key=lambda o: o["placed_at"], reverse=True)
        prefetch = ["These candidates were pre-resolved from the caller's phone number. This is a "
                    "candidate list, not the answer — tag only the order(s) this conversation is "
                    "actually about (match on product, issue, and timing).",
                    f'- Customer "{cust["full_name"]}" (from phone {cust["phone"]}) — recent orders:']
        prefetch += [order_line(o) for o in candidates]
        if rng.random() < 0.5:
            existing = [f"- Customer: {cust['full_name']} (from Contact Resolution (Implicit))"]

    comm_id = "%08x-%04x-%04x-%04x-%012x" % (
        rng.getrandbits(32), rng.getrandbits(16), rng.getrandbits(16),
        rng.getrandbits(16), rng.getrandbits(48))
    sender = cust["phone"] if channel == "Call" else cust["email"]
    parts = ["## Communication", "",
             f"- **ID**: {comm_id}",
             f"- **Channel**: {channel}",
             "- **Direction**: Inbound",
             f"- **Timestamp**: 2026-06-{rng.randint(10, 17):02d}T{rng.randint(8, 19):02d}:"
             f"{rng.randint(0, 59):02d}:00Z",
             f"- **Sender**: {sender} (Name: {cust['full_name']}, Type: External)",
             "- **Recipient**: support@shop.example.com (Name: Support Team, Type: Internal)",
             "", "### Call Transcript" if channel == "Call" else "### Content", content]
    if prefetch:
        parts += ["", "### Pre-Fetched Context", *prefetch]
    if existing:
        parts += ["", "### Existing References", *existing]
    text = "\n".join(parts)

    truth = {
        "classification": {"category": category, "subcategory": sub},
        "primary": primary,
        "required": required,
        "optional": optional,
        "lead": lead,
    }
    return truth, text


# --- Task hooks --------------------------------------------------------------

_rng = random.Random()


def sample() -> Sample:
    seed = _rng.randrange(1, 10**9)
    _, text = generate(seed)
    return Sample(input_id=f"comm-{seed}", text=text)


def lookup_truth(input_id: str) -> dict | None:
    m = re.fullmatch(r"comm-(\d+)", input_id or "")
    if not m:
        return None
    truth, _ = generate(int(m.group(1)))
    return truth


def _extract_json(text: str) -> dict | None:
    blocks = re.findall(r"```(?:json)?\s*\n(.*?)```", text, re.DOTALL)
    candidates = [b.strip() for b in blocks]
    if not candidates:
        m = re.search(r"\{.*\}", text, re.DOTALL)  # last-ditch: first { … last }
        if m:
            candidates = [m.group()]
    for cand in reversed(candidates):
        try:
            obj = json.loads(cand)
            if isinstance(obj, dict):
                return obj
        except json.JSONDecodeError:
            continue
    return None


def parse_output(text: str) -> dict:
    obj = _extract_json(text) or {}
    cls = obj.get("classification") or {}
    category = str(cls.get("category") or "").strip()
    subcategory = str(cls.get("subcategory") or "").strip().lower()
    entities = []
    raw_entities = obj.get("entities")
    if isinstance(raw_entities, list):
        for e in raw_entities:
            if not isinstance(e, dict):
                continue
            fid = e.get("friendly_id")
            eid = e.get("entity_id")
            entities.append({
                "entity_type": str(e.get("entity_type") or "").strip(),
                "friendly_id": str(fid).strip() if fid else None,
                "entity_id": str(eid).strip() if eid else None,
                "confidence": e.get("confidence"),
                "reasoning": str(e.get("reasoning") or "").strip(),
            })
    return {
        "parsed_ok": _extract_json(text) is not None,
        "category": category,
        "subcategory": subcategory,
        "entities": entities,
        "has_entities_list": isinstance(raw_entities, list),
        "search_summary": str(obj.get("search_summary") or "").strip(),
        "raw": text,
    }


def _gold_keys(entities: list[dict]) -> tuple[set[str], set[str]]:
    return ({e["friendly_id"].upper() for e in entities if e.get("friendly_id")},
            {e["entity_id"] for e in entities if e.get("entity_id")})


def _covers(pred: dict, fids: set[str], uuids: set[str]) -> bool:
    return (bool(pred["friendly_id"]) and (pred["friendly_id"] or "").upper() in fids) \
        or (bool(pred["entity_id"]) and pred["entity_id"] in uuids)


def _entity_keys(e: dict) -> tuple[set[str], set[str]]:
    fids = {e["friendly_id"].upper()} if e.get("friendly_id") else set()
    uuids = {e["entity_id"]} if e.get("entity_id") else set()
    return fids, uuids


def _is_covered(gold_entity: dict, preds: list[dict]) -> bool:
    fids, uuids = _entity_keys(gold_entity)
    return any(_covers(p, fids, uuids) for p in preds)


# A fabricated ID is only a hallucination if it is SHAPED like a transaction id / UUID
# yet does not exist — a customer name or product name in friendly_id is not an invented id.
_ID_SHAPE = re.compile(r"(?i)^(ord|shp|rma|sku|wo|tkt)-\w")
_UUID_SHAPE = re.compile(r"(?i)^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$")


def score(truth: dict, parsed: dict) -> dict[str, float]:
    # --- Task 1: classification ---
    category, sub = truth["classification"]["category"], truth["classification"]["subcategory"]
    valid = parsed["category"] in TAXONOMY and parsed["subcategory"] in TAXONOMY.get(parsed["category"], [])
    cat_ok = float(parsed["category"].lower() == category.lower())
    sub_ok = float(parsed["subcategory"] == sub)
    intent_exact = cat_ok * sub_ok

    # --- Task 2: entity tagging ---
    preds = parsed["entities"]
    seen, gold_u = set(), []  # de-dup gold by (friendly_id or uuid)
    for g in [truth["primary"], *truth["required"], *truth["optional"]]:
        key = (g.get("friendly_id") or "").upper() or g["entity_id"]
        if key not in seen:
            seen.add(key)
            gold_u.append(g)
    gold_fids, gold_uuids = _gold_keys(gold_u)
    prim, required = truth["primary"], truth["required"]

    matched_preds = [p for p in preds if _covers(p, gold_fids, gold_uuids)]
    covered_gold = sum(_is_covered(g, preds) for g in gold_u)
    precision = len(matched_preds) / len(preds) if preds else 0.0
    recall = covered_gold / len(gold_u) if gold_u else 0.0
    f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0
    primary_found = float(_is_covered(prim, preds))
    required_recall = (sum(_is_covered(g, preds) for g in required) / len(required)) if required else 0.0
    entity_id_correct = float(any(p["entity_id"] == prim["entity_id"] for p in preds if p["entity_id"]))

    no_hallucinated = 1.0
    for p in preds:
        fid, eid = p["friendly_id"], p["entity_id"]
        if fid and _ID_SHAPE.match(fid) and fid.upper() not in WORLD_FRIENDLY_IDS:
            no_hallucinated = 0.0
        if eid and _UUID_SHAPE.match(eid) and eid not in WORLD_UUIDS:
            no_hallucinated = 0.0

    return {
        "intent_format_ok": float(valid),
        "category_match": cat_ok,
        "subcategory_match": sub_ok,
        "intent_exact": intent_exact,
        "tag_format_ok": float(parsed["has_entities_list"] and bool(parsed["search_summary"])),
        "primary_found": primary_found,
        "entity_id_correct": entity_id_correct,
        "required_recall": float(required_recall),
        "precision": float(precision),
        "recall": float(recall),
        "f1": float(f1),
        "no_hallucinated_id": no_hallucinated,
        "both_correct": intent_exact * primary_found,
    }


def present(parsed: dict, truth: dict | None) -> dict:
    out: dict = {
        "classification": {"category": parsed["category"] or None,
                           "subcategory": parsed["subcategory"] or None},
        "search_summary": parsed["search_summary"],
        "raw": parsed["raw"],
        "parsed_ok": parsed["parsed_ok"],
        "tags": [],
    }
    gold_fids: set[str] = set()
    gold_uuids: set[str] = set()
    prim = None
    if truth:
        gold = [truth["primary"], *truth["required"], *truth["optional"]]
        gold_fids, gold_uuids = _gold_keys(gold)
        prim = truth["primary"]
        out["gold"] = {
            "classification": truth["classification"],
            "primary": prim,
            "required": truth["required"],
            "optional": truth["optional"],
            "lead": truth.get("lead"),
        }
    for p in parsed["entities"]:
        in_gold = _covers(p, gold_fids, gold_uuids) if truth else None
        is_primary = bool(truth) and prim is not None and (
            ((p["friendly_id"] or "").upper() == (prim["friendly_id"] or "x").upper() and p["friendly_id"])
            or p["entity_id"] == prim["entity_id"])
        out["tags"].append({**p, "in_gold": in_gold, "is_primary": bool(is_primary)})
    return out


# Collapsible explainer panels shown verbatim (trusted HTML) above the input on the page.
_OVERVIEW_HTML = """\
<p>
  <strong>What this is.</strong> A demo of an LLM workload for a hypothetical online-shopping
  customer-support operation. Each input is one support case, against which the model executes
  <strong>two tasks</strong>: it <em>classifies the inquiry intent</em> (a category + subcategory,
  e.g. Shipping&nbsp;→&nbsp;late_delivery) and <em>tags the business entities</em> the case refers
  to — the order, shipment, return, customer, or product. Its output is that classification plus a
  list of tagged entities; when the input was dealt here, both are scored automatically against the
  known-correct answer.
</p>

<p>
  <strong>What an input is.</strong> A block of text assembled from a customer message (e.g. a
  refund request or a late-delivery complaint) together with pre-fetched context, message metadata,
  and any entities that (hypothetical) upstream systems already linked to it. Everything is synthetic
  and generated on the fly: each deal draws a random scenario over a fixed, made-up catalog and
  renders it as an email, chat, or call.
</p>

<p>
  <strong>Tools the model can use.</strong> Classification needs no tools. Two tools are available
  for the tagging task:
</p>
<ul>
  <li>
    <code>search_entities</code> — fuzzy text search over the catalog. Used to find an entity by
    name, number, or description (e.g. the order behind a product name, or to confirm an order ID
    quoted in the message).
  </li>
  <li>
    <code>execute_sql</code> — read-only SQL over the order database. Used for what search can't do:
    resolving a customer by phone number, joining orders to shipments, or filtering by date.
  </li>
  <li>
    <strong>The "database"</strong> is a small, fixed synthetic dataset committed in the repo
    (customers, products, orders, order&nbsp;items, shipments, returns). <code>execute_sql</code>
    queries it in-process (via duckdb) and <code>search_entities</code> indexes the same data — so
    every tool result is a real query result over that dataset (deterministic).
  </li>
</ul>

<p>
  <strong>How to use it.</strong> Click <em>Deal me a communication</em> for a fresh input, then
  <em>Run&nbsp;▸</em>. You'll see the classification decision, the tagged entities, the tools the
  model called, and (for dealt inputs) the auto-scores. You can provide feedback using 👍/👎 and the
  note box. You can also provide your own text input instead of dealing one (but there will be no
  ground-truth).
</p>

<p>
  <strong>The traffic.</strong> Every run becomes one trace in Langfuse — the prompt, each tool
  call, the final answer, and all the scores — which is the data this demo exists to produce. To
  see only this demo's traffic, filter the Traces tab to the tag <code>order-comms</code> (every run
  here carries it). Dashboard:
  <a href="https://cloud.langfuse.com/project/cmq8ojtt101fhad0i3qqaxiza/traces" target="_blank" rel="noopener">https://cloud.langfuse.com/project/cmq8ojtt101fhad0i3qqaxiza/traces</a>
</p>
"""

# The walkthrough below uses a REAL dealt input (seed 1) and REAL tool payloads captured
# from the actual tools — edit the prose freely, but the JSON mirrors what the app produces.
_EXAMPLE_HTML = """\
<p>
  A representative run on a real dealt input. The input is an inbound <strong>chat</strong> that
  references order and shipping ids.
</p>

<p class="step">Example input</p>
<pre>## Communication

- Channel: Chat
- Direction: Inbound
- Timestamp: 2026-06-10T14:27:00Z
- Sender: sasha.nakamura21@mailbox.example.com (Name: Sasha Nakamura, Type: External)
- Recipient: support@shop.example.com (Name: Support Team, Type: Internal)

### Content
My plush bear shipment seems lost — no movement for over a week and nothing at my
door. Order number is ORD-260606-4070. Reference SHP-260608-1862.

### Pre-Fetched Context
- SHP-260608-1862 (in message) -> Order ORD-260606-4070 — plush bear, placed
  2026-06-06, status shipped, Shipment SHP-260608-1862 exception</pre>

<p class="step">Step 1 — classify intent (no tool call)</p>
<p>
  From the wording alone the model decides
  <code>{"category": "Shipping", "subcategory": "lost_package"}</code>. Classification never uses a
  tool.
</p>

<p class="step">Step 2 — tag entities, tool call 1: <code>search_entities</code></p>
<p>The model resolves the order number quoted in the message. Arguments it emits:</p>
<pre>{ "search_terms": "ORD-260606-4070", "entity_types": ["Order"],
  "reason": "Resolve the order quoted in the message" }</pre>
<p>Tool result (top hit; nine lower-scored hits omitted):</p>
<pre>{
  "result_count": 10,
  "results": [
    { "id": "46839f5b-b827-91a9-3e05-bbca736619a2", "type": "Order", "score": 109.0,
      "fields": { "order_number": "ORD-260606-4070", "status": "shipped",
                  "customer_name": "Sasha Nakamura", "ship_to": "Ironwood, MI",
                  "total": 178.36 } }
  ]
}</pre>

<p class="step">Step 3 — tag entities, tool call 2: <code>execute_sql</code></p>
<p>
  The order doc has no shipment ID, so the model queries the database for the shipment row (search
  can't return the shipment's UUID). Arguments:
</p>
<pre>{ "sql": "SELECT shipment_id, shipment_number, ship_status, est_delivery
          FROM shop.shipments WHERE shipment_number = 'SHP-260608-1862'",
  "reason": "Get the shipment row and its id" }</pre>
<p>Tool result:</p>
<pre>{
  "row_count": 1,
  "columns": ["shipment_id", "shipment_number", "ship_status", "est_delivery"],
  "rows": [
    { "shipment_id": "75e88d7e-ab67-eeae-e3d7-e9dc4109752a",
      "shipment_number": "SHP-260608-1862", "ship_status": "exception",
      "est_delivery": "2026-06-15 20:59:00" }
  ]
}</pre>

<p class="step">Step 4 — final answer</p>
<p>
  The model combines the two tool results into one JSON object: the lost shipment as the primary
  entity plus its parent order, each carrying the UUID it pulled from a tool result.
</p>
<pre>{
  "classification": { "category": "Shipping", "subcategory": "lost_package" },
  "entities": [
    { "entity_type": "Shipment", "friendly_id": "SHP-260608-1862",
      "entity_id": "75e88d7e-ab67-eeae-e3d7-e9dc4109752a", "confidence": 0.95,
      "reasoning": "The lost shipment named in the message; ship_status 'exception' fits." },
    { "entity_type": "Order", "friendly_id": "ORD-260606-4070",
      "entity_id": "46839f5b-b827-91a9-3e05-bbca736619a2", "confidence": 0.9,
      "reasoning": "Parent order quoted in the message." }
  ],
  "search_summary": "Resolved the order via search_entities; fetched its shipment + id via SQL."
}</pre>
<p>
  Against this input's known-correct answer that scores a clean pass: right category and
  subcategory, the correct primary entity tagged, the right UUIDs supplied (only obtainable from a
  tool result), and no invented IDs.
</p>
"""

# Panel 3 is GENERATED from the live TAXONOMY / TOOLS / committed data and from real tool
# output, so it stays in sync automatically. Edit the prose in _build_details_html() below;
# the label list, tool schemas, payloads, and DB schema are derived — don't hand-copy them.


def _esc(s: object) -> str:
    return str(s).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


def _pre(obj: object) -> str:
    text = obj if isinstance(obj, str) else json.dumps(obj, indent=2, default=str)
    return f"<pre>{_esc(text)}</pre>"


def _build_details_html() -> str:
    schema_sources = [
        ("shop.customers", CUSTOMERS), ("shop.products", PRODUCTS), ("shop.orders", ORDERS),
        ("shop.order_items", ORDER_ITEMS), ("shop.shipments", SHIPMENTS), ("shop.returns", RETURNS),
    ]
    n_labels = sum(len(v) for v in TAXONOMY.values())
    tax_items = "\n".join(
        f"  <li><strong>{_esc(cat)}</strong> — {_esc(', '.join(subs))}</li>"
        for cat, subs in TAXONOMY.items())

    fields_by_type: dict[str, list[str]] = {}
    for doc in INDEX:  # insertion order: Order, Shipment, Return, Product, Customer
        fields_by_type.setdefault(doc["type"], list(doc["fields"].keys()))
    field_items = "\n".join(
        f"  <li><strong>{_esc(t)}</strong>: {_esc(', '.join(keys))}</li>"
        for t, keys in fields_by_type.items())

    schema_items = "\n".join(
        f"  <li><code>{name}</code>({_esc(', '.join(src[0].keys()))})</li>"
        for name, src in schema_sources)

    # Real payloads, captured from the actual tools so the documentation can't drift.
    search_ok = json.loads(search_entities(ORDERS[0]["order_number"], ["Order"]))
    search_ok["results"] = search_ok["results"][:1]  # trim to the top hit for brevity
    search_short = json.loads(search_entities("ab"))
    search_none = json.loads(search_entities("zzqqxx yyzzwwvv"))
    sql_ok = json.loads(execute_sql(
        "SELECT order_number, order_status, placed_at FROM shop.orders LIMIT 2", "example"))
    sql_err = json.loads(execute_sql("DELETE FROM shop.orders", "example"))

    return f"""\
<p class="step">Classification label space (Task&nbsp;1)</p>
<p>
  The model outputs one <code>{{"category", "subcategory"}}</code> pair from this fixed taxonomy
  ({n_labels} labels in all). Each category's <code>general</code> subcategory is the fallback for
  when the specific one can't be pinned down.
</p>
<ul>
{tax_items}
</ul>

<p class="step">Tool definitions (Task 2)</p>
{_pre(TOOLS)}

<p class="step">search_entities — result payloads</p>
<p>
  Success returns a flat, score-sorted list (exact friendly-ID matches are boosted). The
  <code>fields</code> object differs by entity type. Top hit for a full order-number search:
</p>
{_pre(search_ok)}
<p>The <code>fields</code> extracted for each entity type:</p>
<ul>
{field_items}
</ul>
<p>Errors come back as a result (not raised) — e.g. a term under 3 characters, or no match at all:</p>
{_pre(search_short)}
{_pre(search_none)}

<p class="step">execute_sql — result payloads</p>
<p>
  Read-only (only <code>SELECT</code> / <code>WITH</code> run); a missing <code>LIMIT</code> is
  capped to 50, and at most 50 rows are returned:
</p>
{_pre(sql_ok)}
<p>Anything that isn't a read-only query comes back as an error result:</p>
{_pre(sql_err)}

<p class="step">Database schema (the <code>shop</code> tables)</p>
<p>
  Both tools read this one committed, synthetic dataset. An entity's <code>entity_id</code> in the
  final answer is its UUID from these tables.
</p>
<ul>
{schema_items}
</ul>
<p>
  Note: <code>customers.phone</code> is deliberately <em>not</em> indexed by search — resolving a
  caller by phone number has to go through <code>execute_sql</code>.
</p>
"""


_DETAILS_HTML = _build_details_html()

PANELS = [
    {"title": "What is this?", "html": _OVERVIEW_HTML},
    {"title": "Walk through an example run", "html": _EXAMPLE_HTML},
    {"title": "Tasks, tools, and DB details", "html": _DETAILS_HTML},
]

TASK = Task(
    order=0,  # first in the nav — the headline demo
    id="order-comms",
    title="Support Comms",
    tagline="Classify a customer message and tag the orders and entities it's about",
    system_prompt=SYSTEM_PROMPT,
    ui={
        "output": "comms",
        "input_label": "Support case (customer message + pre-fetched context + metadata)",
        "placeholder": "Deal a communication, or paste your own email/chat/call transcript...",
        "deal_label": "Deal me a communication",
        "panels": PANELS,
    },
    sample=sample,
    lookup_truth=lookup_truth,
    parse_output=parse_output,
    score=score,
    present=present,
    tools=TOOLS,
    execute_tool=execute_tool,
)