Update app.py

app.py

@@ -1,378 +1,340 @@
 """
-PlotWeaver Voice Agent — HuggingFace Space (Gradio 6 + ZeroGPU)
-================================================================
-Hausa-first conversational AI for African banks, telecoms, and delivery.
-Pipeline:
-  ASR  (openai/whisper-small)            ── on GPU (@spaces.GPU)
-    → NLU (embedding similarity, nlu.py) ── on CPU (~200ms; GPU would be slower)
-    → Dialogue FSM (dialogue.py) + Backend (backend.py)
-    → TTS  (facebook/mms-tts-hau)        ── on GPU (@spaces.GPU)
-
-ZeroGPU notes
--------------
-* Requires a PRO account + ZeroGPU hardware selected in Space settings.
-* Models are placed on cuda at MODULE level (the ZeroGPU contract), even though
-  a real GPU is only attached inside @spaces.GPU functions.
-* Decorated functions return CPU/numpy data (never live CUDA tensors).
-* Low `duration` values raise queue priority — ASR/TTS per turn are seconds.
-
-Secrets / variables (Space → Settings):
-  HF_TOKEN         write token, for log syncing (see logging_util.py)
-  PW_LOG_DATASET   e.g. "plotweaver/voice-agent-logs"
+NLU Module — Multi-Intent Decomposition + Entity Extraction
+=============================================================
+Fixes P0/P1 from the feedback:
+
+  - Decomposes ONE user message into a LIST of tasks (compound requests)
+  - Extracts ALL entities present in the message (slot prefill —
+    "send money to abu" never re-asks for the recipient)
+  - Returns per-task confidence so destructive intents can be gated
+  - Distinguishes "ask about X" from "do X" (branch location ≠ block card)
+
+Backend chain (first available wins):
+  1. LLM_API    — HF Serverless Inference (set HF_TOKEN) — best quality
+  2. LLM_LOCAL  — Qwen2.5-1.5B-Instruct loaded in-process — good, slower
+  3. RULES      — improved keyword rules — degraded but never crashes
+
+All backends return the same schema:
+
+  {
+    "tasks": [
+      {
+        "intent": "send_money",
+        "confidence": 0.93,
+        "slots": {"recipient": "abu", "amount": "350000"},
+        "utterance_span": "send 350000 to abu"
+      },
+      ...
+    ],
+    "backend": "llm_api"
+  }
 """
-from __future__ import annotations
+
 import os
-import time
-import uuid
-import html as html_lib
+import re
+import json
+import logging
 from typing import Optional
 
-import gradio as gr
-import numpy as np
-import torch
-import spaces
-from transformers import (
-    VitsModel, AutoTokenizer,
-    WhisperProcessor, WhisperForConditionalGeneration,
-)
-
-from dialogue import (
-    DialogueState, SCENARIOS,
-    render_prompt, get_expected_slot, transition,
-)
-from nlu import parse as nlu_parse
-from logging_util import log_turn
-
-
-# ---------------------------------------------------------------------------
-# Device selection — cuda at module level on ZeroGPU, else local CPU/GPU
-# ---------------------------------------------------------------------------
-_IS_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))
-if _IS_ZEROGPU or torch.cuda.is_available():
-    DEVICE = "cuda"
-else:
-    DEVICE = "cpu"
-
-
-# ---------------------------------------------------------------------------
-# Model loading (module level — required by ZeroGPU)
-# ---------------------------------------------------------------------------
-print(f"Loading models on {DEVICE}…")
-
-_asr_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
-_asr_model = (
-    WhisperForConditionalGeneration
-    .from_pretrained("openai/whisper-small")
-    .to(DEVICE)
-    .eval()
-)
-
-_tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-hau")
-_tts_model = VitsModel.from_pretrained("facebook/mms-tts-hau").to(DEVICE).eval()
-# PlotWeaver Hausa production voice config (tune here, not at call sites).
-_tts_model.speaking_rate = 0.85
-_tts_model.noise_scale = 0.4
-_TTS_SEED = 42
-
-print("Models ready.")
-
-
-@spaces.GPU(duration=15)
-def transcribe_hausa(audio_tuple) -> str:
-    if audio_tuple is None:
-        return ""
-    sample_rate, audio_array = audio_tuple
-    if audio_array is None or len(audio_array) == 0:
-        return ""
-    if audio_array.dtype != np.float32:
-        audio_array = audio_array.astype(np.float32) / np.iinfo(audio_array.dtype).max
-    if audio_array.ndim > 1:
-        audio_array = audio_array.mean(axis=1)
-    # Cap at 30s (Whisper training chunk size)
-    max_samples = sample_rate * 30
-    if len(audio_array) > max_samples:
-        audio_array = audio_array[:max_samples]
-    if sample_rate != 16000:
-        import scipy.signal
-        num_samples = int(len(audio_array) * 16000 / sample_rate)
-        audio_array = scipy.signal.resample(audio_array, num_samples).astype(np.float32)
-
-    inputs = _asr_processor(audio_array, sampling_rate=16000, return_tensors="pt")
-    forced_ids = _asr_processor.get_decoder_prompt_ids(language="hausa", task="transcribe")
-    feats = inputs.input_features.to(DEVICE)
-    with torch.no_grad():
-        ids = _asr_model.generate(feats, forced_decoder_ids=forced_ids, max_new_tokens=128)
-    text = _asr_processor.batch_decode(ids, skip_special_tokens=True)[0].strip()
-    return text
-
-
-@spaces.GPU(duration=10)
-def synthesize_hausa(text: str) -> Optional[tuple]:
-    if not text.strip():
-        return None
-    torch.manual_seed(_TTS_SEED)  # reproducible voice (prod parity)
-    inputs = _tts_tokenizer(text, return_tensors="pt")
-    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
-    with torch.no_grad():
-        out = _tts_model(**inputs).waveform
-    audio = out.squeeze().cpu().numpy().astype(np.float32)
-    return (_tts_model.config.sampling_rate, audio)
-
-
-# ---------------------------------------------------------------------------
-# WhatsApp-style HTML rendering
-# ---------------------------------------------------------------------------
-def _now() -> str:
-    return time.strftime("%H:%M")
-
-
-def _user_bubble(text: str, is_voice: bool) -> str:
-    t = html_lib.escape(text)
-    if is_voice:
-        bars = "".join(
-            f'<span style="height:{4 + int(8 * abs(np.sin(i * 0.7)))}px;"></span>'
-            for i in range(20)
-        )
-        return f'''<div class="pw-b user">
-  <div class="pw-voice-row">
-    <div class="pw-voice-icon">▶</div>
-    <div class="pw-voice-bars">{bars}</div>
-  </div>
-  <div style="font-size:12px;color:#667781;margin-top:3px;">"{t}"</div>
-  <div class="pw-b-meta">{_now()} ✓✓</div>
-</div>'''
-    return f'<div class="pw-b user">{t}<div class="pw-b-meta">{_now()} ✓✓</div></div>'
-
-
-def _bot_bubble(text_ha: str, text_en: str) -> str:
-    ha = html_lib.escape(text_ha)
-    en = html_lib.escape(text_en)
-    return f'''<div class="pw-b bot">
-  <div>{ha}</div>
-  <div class="pw-b-trans">{en}</div>
-  <div class="pw-b-meta">{_now()} ✓✓</div>
-</div>'''
-
-
-def render_whatsapp(session: dict) -> str:
-    vertical = session.get("vertical", "bank") if session else "bank"
-    name = SCENARIOS[vertical]["name"]
-    avatar = {"bank": "PB", "telecom": "PT", "ecommerce": "PD"}[vertical]
-    escalated = session.get("escalate_to_human", False) if session else False
-
-    bubbles = []
-    for msg in session.get("history", []) if session else []:
-        if msg["role"] == "user":
-            bubbles.append(_user_bubble(msg["text"], msg.get("is_voice", False)))
+logger = logging.getLogger(__name__)
+
+# ── Intent catalogue (shared by all backends) ────────────────────────────────
+INTENT_SCHEMA = {
+    "greeting":         {"slots": [],                          "destructive": False},
+    "balance_inquiry":  {"slots": ["account_id"],              "destructive": False},
+    "send_money":       {"slots": ["recipient", "amount", "account_id"], "destructive": True},
+    "bill_payment":     {"slots": ["account_id", "amount"],    "destructive": True},
+    "block_card":       {"slots": ["account_id"],              "destructive": True},
+    "branch_info":      {"slots": [],                          "destructive": False},
+    "card_request":     {"slots": [],                          "destructive": False},
+    "report_issue":     {"slots": ["issue_desc"],              "destructive": False},
+    "track_order":      {"slots": ["order_id"],                "destructive": False},
+    "return_item":      {"slots": ["order_id", "return_reason"], "destructive": False},
+    "human_agent":      {"slots": [],                          "destructive": False},
+    "confirmation_yes": {"slots": [],                          "destructive": False},
+    "confirmation_no":  {"slots": [],                          "destructive": False},
+    "cancel":           {"slots": [],                          "destructive": False},
+    "goodbye":          {"slots": [],                          "destructive": False},
+    "unknown":          {"slots": [],                          "destructive": False},
+}
+
+NLU_SYSTEM_PROMPT = """You are the NLU module of a customer-service voice agent.
+Decompose the user's message into ALL tasks it contains, in order.
+Extract every entity present. Never invent entities that are not in the text.
+
+Intents: greeting, balance_inquiry, send_money, bill_payment, block_card,
+branch_info, card_request, report_issue, track_order, return_item,
+human_agent, confirmation_yes, confirmation_no, cancel, goodbye, unknown.
+
+Slots: recipient, amount, account_id, location, issue_desc, order_id, return_reason.
+
+CRITICAL disambiguation rules:
+- "where is your branch so I can get my card" = branch_info + card_request.
+  It is NOT block_card. Only choose block_card if the user explicitly wants to
+  BLOCK, FREEZE, or DEACTIVATE a card.
+- A message can contain multiple tasks joined by "and", "also", "then".
+  Output one task per action. "check my balance and send 5000 to musa"
+  = [balance_inquiry, send_money{recipient: musa, amount: 5000}].
+- If the user answers a question (e.g. gives a reason like "too small"),
+  map it to the slot of the pending task, intent = the pending intent.
+- Confidence in [0,1]: how sure you are of the INTENT (not the slots).
+
+Respond with ONLY valid JSON, no markdown, no commentary:
+{"tasks":[{"intent":"...","confidence":0.0,"slots":{},"utterance_span":"..."}]}"""
+
+
+class NLU:
+
+    def __init__(self, prefer: str = "auto"):
+        self.hf_token = os.getenv("HF_TOKEN", "")
+        self.api_model = os.getenv(
+            "NLU_API_MODEL", "Qwen/Qwen2.5-72B-Instruct")
+        self.local_model_id = os.getenv(
+            "NLU_LOCAL_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
+        self._local_pipe = None
+        self.prefer = prefer
+
+    # ── Public API ────────────────────────────────────────────────────────────
+
+    def parse(self, text: str, pending_intent: Optional[str] = None,
+              pending_slot: Optional[str] = None) -> dict:
+        """
+        text           : English pivot text of the user turn
+        pending_intent : intent currently awaiting a slot (context for the LLM)
+        pending_slot   : which slot we asked for last turn
+        """
+        context = ""
+        if pending_intent and pending_slot:
+            context = (f"\nContext: you previously asked the user for the "
+                       f"'{pending_slot}' of a '{pending_intent}' task. "
+                       f"A short answer likely fills that slot.")
+
+        for backend in self._backend_order():
+            try:
+                result = backend(text, context)
+                if result and result.get("tasks"):
+                    result = self._sanitize(result)
+                    logger.info(f"NLU[{result['backend']}]: "
+                                f"{json.dumps(result['tasks'])[:200]}")
+                    return result
+            except Exception as e:
+                logger.warning(f"NLU backend failed ({backend.__name__}): {e}")
+        # Absolute last resort
+        return {"tasks": [{"intent": "unknown", "confidence": 0.0,
+                            "slots": {}, "utterance_span": text}],
+                "backend": "none"}
+
+    # ── Backend chain ─────────────────────────────────────────────────────────
+
+    def _backend_order(self):
+        if self.prefer == "rules":
+            return [self._rules_backend]
+        chain = []
+        if self.hf_token:
+            chain.append(self._api_backend)
+        chain.append(self._local_backend)
+        chain.append(self._rules_backend)
+        return chain
+
+    # ── 1. HF Serverless Inference API ───────────────────────────────────────
+
+    def _api_backend(self, text: str, context: str) -> Optional[dict]:
+        import requests
+        url = f"https://api-inference.huggingface.co/models/{self.api_model}/v1/chat/completions"
+        payload = {
+            "model": self.api_model,
+            "messages": [
+                {"role": "system", "content": NLU_SYSTEM_PROMPT + context},
+                {"role": "user",   "content": text},
+            ],
+            "max_tokens": 400,
+            "temperature": 0.1,
+        }
+        r = requests.post(url, json=payload, timeout=20,
+                          headers={"Authorization": f"Bearer {self.hf_token}"})
+        r.raise_for_status()
+        raw = r.json()["choices"][0]["message"]["content"]
+        parsed = self._extract_json(raw)
+        if parsed:
+            parsed["backend"] = "llm_api"
+        return parsed
+
+    # ── 2. Local small LLM ────────────────────────────────────────────────────
+
+    def _local_backend(self, text: str, context: str) -> Optional[dict]:
+        if self._local_pipe is None:
+            logger.info(f"Loading local NLU model {self.local_model_id} …")
+            from transformers import pipeline as hf_pipeline
+            import torch
+            self._local_pipe = hf_pipeline(
+                "text-generation",
+                model=self.local_model_id,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                device_map="auto",
+            )
+        messages = [
+            {"role": "system", "content": NLU_SYSTEM_PROMPT + context},
+            {"role": "user",   "content": text},
+        ]
+        out = self._local_pipe(messages, max_new_tokens=400,
+                                do_sample=False, temperature=None, top_p=None)
+        raw = out[0]["generated_text"][-1]["content"]
+        parsed = self._extract_json(raw)
+        if parsed:
+            parsed["backend"] = "llm_local"
+        return parsed
+
+    # ── 3. Improved rules (never fails) ───────────────────────────────────────
+
+    def _rules_backend(self, text: str, context: str) -> dict:
+        """
+        Better than the old FSM keywords:
+          - splits on conjunctions to find MULTIPLE tasks
+          - extracts entities per clause
+          - branch_info vs block_card disambiguation
+        """
+        t = text.lower().strip()
+
+        # Split compound message into clauses
+        clauses = re.split(r'\b(?:and also|and then|then|and|also|;|\. )\b', t)
+        clauses = [c.strip() for c in clauses if c.strip()]
+
+        tasks = []
+        for clause in clauses:
+            task = self._rules_classify_clause(clause)
+            if task:
+                tasks.append(task)
+
+        # Merge duplicate consecutive intents (e.g. "and" split an entity off)
+        merged = []
+        for task in tasks:
+            if merged and merged[-1]["intent"] == task["intent"]:
+                merged[-1]["slots"].update(task["slots"])
+                merged[-1]["utterance_span"] += " " + task["utterance_span"]
+            else:
+                merged.append(task)
+
+        if not merged:
+            merged = [{"intent": "unknown", "confidence": 0.3,
+                       "slots": {}, "utterance_span": t}]
+
+        return {"tasks": merged, "backend": "rules"}
+
+    def _rules_classify_clause(self, clause: str) -> Optional[dict]:
+        slots = {}
+
+        # ── Entity extraction (always, regardless of intent) ────────────────
+        # In a money-action clause ("send/transfer/pay X to Y"), the number is
+        # an AMOUNT. Only treat 6-12 digit numbers as account_id when the
+        # clause is about the account itself, or there is no money verb.
+        money_verb = any(v in clause for v in ("send", "transfer", "pay"))
+        account_ctx = any(v in clause for v in ("account", "acct", "number is"))
+        numbers = re.findall(r'\b\d[\d,\.]*\b', clause)
+        for num in numbers:
+            digits = num.replace(",", "").replace(".", "")
+            if money_verb and "amount" not in slots and len(digits) <= 7:
+                slots["amount"] = digits
+            elif (account_ctx or not money_verb) and 6 <= len(digits) <= 12 \
+                    and "account_id" not in slots:
+                slots["account_id"] = digits
+            elif "amount" not in slots and len(digits) <= 7:
+                slots["amount"] = digits
+        # recipient: "to <name>" — take the LAST valid match, skipping verbs
+        # ("I want to send money to abu" must yield 'abu', not 'send')
+        RECIPIENT_STOPWORDS = {
+            "my", "the", "a", "an", "me", "you", "check", "send", "transfer",
+            "pay", "get", "make", "do", "know", "see", "block", "return",
+            "track", "him", "her", "them", "it", "confirm", "cancel"}
+        for m in re.finditer(r'\bto\s+([a-z]{2,20})\b', clause):
+            name = m.group(1)
+            if name not in RECIPIENT_STOPWORDS:
+                slots["recipient"] = name
+        # order id
+        m = re.search(r'\border\s*#?\s*([a-z0-9\-]{4,20})\b', clause)
+        if m:
+            slots["order_id"] = m.group(1)
+
+        # ── Intent (order matters: destructive intents need explicit verbs) ──
+        def has(*kws):
+            return any(kw in clause for kw in kws)
+
+        # branch/location questions BEFORE block_card — fixes P0 #2
+        if has("branch", "closest", "nearest", "location", "where is", "address"):
+            intent, conf = "branch_info", 0.85
+            if has("card", "atm"):
+                # compound: they also want a card — but NOT to block it
+                return {"intent": "branch_info", "confidence": 0.85,
+                        "slots": slots, "utterance_span": clause}
+        elif has("block my card", "block card", "freeze", "deactivate", "stolen", "lost my card"):
+            intent, conf = "block_card", 0.8
+        elif has("send", "transfer") and (slots.get("recipient") or slots.get("amount")):
+            intent, conf = "send_money", 0.85
+        elif has("send money", "transfer money"):
+            intent, conf = "send_money", 0.75
+        elif has("balance", "how much", "asusun"):
+            intent, conf = "balance_inquiry", 0.85
+        elif has("pay", "bill", "recharge", "invoice"):
+            intent, conf = "bill_payment", 0.75
+        elif has("track", "where is my order", "delivery", "shipment"):
+            intent, conf = "track_order", 0.8
+        elif has("return", "refund", "send back"):
+            intent, conf = "return_item", 0.8
+        elif has("problem", "issue", "complaint", "not working", "error"):
+            intent, conf = "report_issue", 0.7
+            slots["issue_desc"] = clause
+        elif has("human", "agent", "person", "operator", "representative"):
+            intent, conf = "human_agent", 0.9
+        elif has("yes", "yep", "correct", "confirm", "sure", "okay", "ok"):
+            intent, conf = "confirmation_yes", 0.8
+        elif has("no", "nope", "wrong", "cancel that"):
+            intent, conf = "confirmation_no", 0.8
+        elif has("hello", "hi ", "good morning", "sannu", "salam"):
+            intent, conf = "greeting", 0.9
+        elif has("bye", "goodbye", "thank"):
+            intent, conf = "goodbye", 0.85
         else:
-            bubbles.append(_bot_bubble(msg.get("text_ha", ""), msg.get("text_en", "")))
-
-    banner = ('<div class="pw-esc-banner">Session escalated to human agent</div>'
-              if escalated else "")
-
-    if not bubbles:
-        body = '<div style="text-align:center;color:#667781;font-size:12px;padding:40px 0;">Send a message to begin…</div>'
-    else:
-        body = "".join(bubbles)
-
-    return f"""
-<div class="pw-phone">
-  <div class="pw-ph-header">
-    <div class="pw-ph-avatar">{avatar}</div>
-    <div>
-      <div class="pw-ph-name">{html_lib.escape(name)}</div>
-      <div class="pw-ph-status">online • voice agent</div>
-    </div>
-  </div>
-  <div class="pw-ph-messages">
-    {banner}
-    {body}
-  </div>
-</div>
-<style>
-.pw-phone {{ max-width: 480px; margin: 0 auto; background: #ECE5DD; border-radius: 14px; overflow: hidden; border: 1px solid #ccc; display: flex; flex-direction: column; min-height: 540px; font-family: -apple-system, "Segoe UI", Roboto, sans-serif; }}
-.pw-ph-header {{ background: #075E54; color: #fff; padding: 10px 14px; display: flex; align-items: center; gap: 10px; }}
-.pw-ph-avatar {{ width: 36px; height: 36px; border-radius: 50%; background: #128C7E; display: flex; align-items: center; justify-content: center; font-weight: 500; font-size: 13px; color: #fff; }}
-.pw-ph-name {{ font-size: 14px; font-weight: 500; line-height: 1.2; }}
-.pw-ph-status {{ font-size: 11px; color: #D4EDE8; }}
-.pw-ph-messages {{ flex: 1; padding: 14px 10px; background: #ECE5DD; background-image: radial-gradient(#D8CFC2 1px, transparent 1px); background-size: 18px 18px; max-height: 480px; overflow-y: auto; min-height: 420px; }}
-.pw-b {{ max-width: 80%; padding: 7px 10px 5px; border-radius: 8px; margin-bottom: 6px; font-size: 13.5px; line-height: 1.4; color: #1f2d1f; word-wrap: break-word; }}
-.pw-b.user {{ background: #DCF8C6; margin-left: auto; border-bottom-right-radius: 2px; }}
-.pw-b.bot {{ background: #fff; margin-right: auto; border-bottom-left-radius: 2px; }}
-.pw-b-meta {{ font-size: 10px; color: #667781; margin-top: 3px; text-align: right; }}
-.pw-b-trans {{ font-size: 11px; color: #667781; font-style: italic; margin-top: 3px; border-top: 1px solid #E5E5E5; padding-top: 3px; }}
-.pw-voice-row {{ display: flex; align-items: center; gap: 8px; }}
-.pw-voice-icon {{ width: 22px; height: 22px; border-radius: 50%; background: #128C7E; color: #fff; font-size: 10px; display: flex; align-items: center; justify-content: center; }}
-.pw-voice-bars {{ flex: 1; height: 14px; display: flex; align-items: center; gap: 2px; }}
-.pw-voice-bars span {{ flex: 1; background: #8D9A9F; border-radius: 1px; }}
-.pw-esc-banner {{ background: #FAEEDA; color: #854F0B; font-size: 12px; padding: 8px 12px; border-radius: 8px; margin-bottom: 10px; border: 1px solid #EF9F27; text-align: center; }}
-</style>
-"""
-
-
-# ---------------------------------------------------------------------------
-# Core turn handler
-# ---------------------------------------------------------------------------
-def run_turn(user_text: str, session: dict, is_voice: bool = False):
-    """Returns (updated_session_dict, bot_audio)."""
-    t0 = time.time()
-    state = DialogueState.from_dict(session) if session else None
-    if state is None:
-        state = DialogueState(session_id="sess_" + uuid.uuid4().hex[:8], vertical="bank")
-
-    state_from = state.current_state
-    expected = get_expected_slot(state.vertical, state.current_state)
-    intent, entities, source = nlu_parse(user_text, expected)
-    state = transition(state, intent, entities)
-
-    prompt = render_prompt(state.vertical, state.current_state, state.slots)
-
-    state.history.append({"role": "user", "text": user_text, "is_voice": is_voice})
-    state.history.append({"role": "bot", "text_ha": prompt["ha"], "text_en": prompt["en"]})
-
-    try:
-        audio = synthesize_hausa(prompt["ha"])
-    except Exception as e:
-        print(f"TTS failed: {e}")
-        audio = None
-
-    log_turn({
-        "session_id": state.session_id,
-        "vertical": state.vertical,
-        "turn": state.turn_count,
-        "is_voice": is_voice,
-        "asr_text": user_text,
-        "expected_slot": expected,
-        "intent": intent,
-        "nlu_source": source,
-        "state_from": state_from,
-        "state_to": state.current_state,
-        "escalated": state.escalate_to_human,
-        "latency_ms": round((time.time() - t0) * 1000),
-    })
-
-    return state.to_dict(), audio
-
-
-# ---------------------------------------------------------------------------
-# Gradio event handlers
-# ---------------------------------------------------------------------------
-def on_vertical_change(vertical: str):
-    state = DialogueState(session_id="sess_" + uuid.uuid4().hex[:8], vertical=vertical)
-    greet = render_prompt(vertical, "greeting", state.slots)
-    state.history.append({"role": "bot", "text_ha": greet["ha"], "text_en": greet["en"]})
-    session = state.to_dict()
-    return session, render_whatsapp(session), None
-
-
-def on_text_submit(text: str, session: dict):
-    if not text or not text.strip():
-        return session, render_whatsapp(session), None, ""
-    new_session, audio = run_turn(text, session, is_voice=False)
-    return new_session, render_whatsapp(new_session), audio, ""
-
-
-def on_audio_submit(audio_data, session: dict):
-    if audio_data is None:
-        return session, render_whatsapp(session), None
-    try:
-        text = transcribe_hausa(audio_data)
-    except Exception as e:
-        print(f"ASR failed: {e}")
-        return session, render_whatsapp(session), None
-    if not text:
-        return session, render_whatsapp(session), None
-    new_session, audio = run_turn(text, session, is_voice=True)
-    return new_session, render_whatsapp(new_session), audio
-
-
-def on_reset(session: dict):
-    vertical = session.get("vertical", "bank") if session else "bank"
-    return on_vertical_change(vertical)
-
-
-# ---------------------------------------------------------------------------
-# Gradio UI (chat-only, minimal components)
-# ---------------------------------------------------------------------------
-CUSTOM_CSS = """
-.gradio-container { max-width: 720px !important; margin: 0 auto !important; }
-#whatsapp-container { padding: 20px 0; }
-"""
-
-with gr.Blocks(css=CUSTOM_CSS, title="PlotWeaver Voice Agent") as demo:
-    gr.HTML("""
-    <div style="text-align:center; padding: 0 0 12px;">
-      <h1 style="margin:0 0 4px; font-size: 22px; font-weight: 500;">PlotWeaver Voice Agent</h1>
-      <p style="margin:0; color: #5f5e5a; font-size: 14px;">Hausa-first conversational AI — pick a vertical, type or speak in Hausa</p>
-    </div>
-    """)
-
-    session_state = gr.State({})
-
-    vertical_radio = gr.Radio(
-        choices=[("PlotWeaver Bank", "bank"),
-                 ("PlotWeaver Telecom", "telecom"),
-                 ("PlotWeaver Delivery", "ecommerce")],
-        value="bank",
-        label="Vertical",
-        container=False,
-    )
-
-    whatsapp_html = gr.HTML(elem_id="whatsapp-container")
-
-    with gr.Row():
-        text_input = gr.Textbox(
-            placeholder="Type in Hausa… e.g. 'duba ma'auni'",
-            label="",
-            scale=4,
-            container=False,
-        )
-        send_btn = gr.Button("Send", scale=1, variant="primary")
-        reset_btn = gr.Button("Reset", scale=1)
-
-    audio_input = gr.Audio(
-        sources=["microphone", "upload"],
-        type="numpy",
-        label="Record or upload Hausa audio (click Stop when done recording)",
-    )
-
-    bot_audio = gr.Audio(
-        label="Bot response (Hausa TTS)",
-        autoplay=True,
-        interactive=False,
-    )
-
-    # Events
-    demo.load(
-        fn=lambda: on_vertical_change("bank"),
-        outputs=[session_state, whatsapp_html, bot_audio],
-    )
-    vertical_radio.change(
-        fn=on_vertical_change,
-        inputs=[vertical_radio],
-        outputs=[session_state, whatsapp_html, bot_audio],
-    )
-    send_btn.click(
-        fn=on_text_submit,
-        inputs=[text_input, session_state],
-        outputs=[session_state, whatsapp_html, bot_audio, text_input],
-    )
-    text_input.submit(
-        fn=on_text_submit,
-        inputs=[text_input, session_state],
-        outputs=[session_state, whatsapp_html, bot_audio, text_input],
-    )
-    audio_input.stop_recording(
-        fn=on_audio_submit,
-        inputs=[audio_input, session_state],
-        outputs=[session_state, whatsapp_html, bot_audio],
-    )
-    reset_btn.click(
-        fn=on_reset,
-        inputs=[session_state],
-        outputs=[session_state, whatsapp_html, bot_audio],
-    )
-
-
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+            return {"intent": "unknown", "confidence": 0.3,
+                    "slots": slots, "utterance_span": clause}
+
+        return {"intent": intent, "confidence": conf,
+                "slots": slots, "utterance_span": clause}
+
+    # ── Helpers ───────────────────────────────────────────────────────────────
+
+    @staticmethod
+    def _extract_json(raw: str) -> Optional[dict]:
+        """Robustly pull the first JSON object out of LLM output."""
+        raw = raw.strip()
+        raw = re.sub(r'^```(?:json)?|```$', '', raw, flags=re.MULTILINE).strip()
+        # find first { … matching last }
+        start = raw.find("{")
+        end   = raw.rfind("}")
+        if start == -1 or end == -1:
+            return None
+        try:
+            return json.loads(raw[start:end + 1])
+        except json.JSONDecodeError:
+            return None
+
+    @staticmethod
+    def _sanitize(result: dict) -> dict:
+        """Validate schema, clamp confidence, drop hallucinated slots."""
+        valid_slots = {"recipient", "amount", "account_id", "location",
+                       "issue_desc", "order_id", "return_reason"}
+        clean_tasks = []
+        for task in result.get("tasks", []):
+            intent = task.get("intent", "unknown")
+            if intent not in INTENT_SCHEMA:
+                intent = "unknown"
+            conf = float(task.get("confidence", 0.5))
+            conf = max(0.0, min(1.0, conf))
+            slots = {k: str(v).strip() for k, v in (task.get("slots") or {}).items()
+                     if k in valid_slots and v not in (None, "", "null", "None")}
+            clean_tasks.append({
+                "intent": intent, "confidence": conf, "slots": slots,
+                "utterance_span": str(task.get("utterance_span", ""))[:200],
+            })
+        result["tasks"] = clean_tasks or [
+            {"intent": "unknown", "confidence": 0.0, "slots": {},
+             "utterance_span": ""}]
+        return result