app.py

"""
Commerce Operations Agent — demo Space for SkyAsl/Qwen3.5-9B-com-agent.

The adapter is a structured-JSON transducer over two fixed tasks
(capability_advice, operation_plan). This app reproduces the exact training
contract — same system prompts, same compact-JSON user payloads — and
validates every response against the guardrails the model was trained on.
"""

import json
import os

import gradio as gr
import torch
from peft import PeftModel
from transformers import AutoModelForImageTextToText, AutoTokenizer

ZERO_GPU = False
try:  # ZeroGPU is only present on Spaces GPU hardware
    import spaces

    GPU = spaces.GPU(duration=120)
    ZERO_GPU = True
    print("ZeroGPU active — GPU will be allocated per request.")
except Exception as _e:  # local / CPU fallback — decorator becomes a no-op
    # Loud on purpose: on ZeroGPU hardware this path means no GPU is ever
    # allocated, and a 9B model will fail or crawl on CPU.
    print(f"WARNING: `spaces` unavailable ({_e}). Running WITHOUT ZeroGPU.")

    def GPU(fn):
        return fn


BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3.5-9B")
ADAPTER = os.environ.get("ADAPTER", "SkyAsl/Qwen3.5-9B-com-agent")

# Exact system prompts from the fine-tuning data. Do not paraphrase — the
# adapter keys off these strings.
SYSTEM = {
    "capability_advice": (
        "TASK=capability_advice. Recommend only registered operations. "
        "Return valid JSON."
    ),
    "operation_plan": (
        "TASK=operation_plan. Select only supplied candidate IDs. "
        "Do not execute operations. Return valid JSON."
    ),
}

# --- presets, taken verbatim from the held-out test split -----------------

PRESET_CAPABILITY = {
    "available_operations": ["restock_inventory", "flag_issue", "draft_order_followup"],
    "store_summary": {
        "low_stock": 15,
        "products": 460,
        "stockout_risks": 4,
        "supplier_coverage": 0.47,
        "unfulfilled_orders": 14,
    },
}

PRESET_PLAN_ACT = {
    "candidates": [
        {
            "candidate_id": "restock:063-00-0",
            "operation_type": "restock_inventory",
            "facts": {
                "committed": 0,
                "daily_velocity": 4.4,
                "data_complete": True,
                "eligible": True,
                "inbound": 1,
                "lead_time_days": 12,
                "minimum_order_quantity": 20,
                "net_available": 54,
                "on_hand": 53,
                "pack_size": 12,
                "reorder_point": 61,
                "reorder_quantity": 132,
                "safety_stock": 8,
                "sales_lookback_days": 45,
                "sales_units": 198,
                "supplier_configured": True,
                "target_coverage_days": 38,
                "target_stock": 176,
                "validated_quantity": 132,
            },
        }
    ],
    "enabled_operations": ["restock_inventory"],
}

# Same candidate, but the operation is NOT enabled. Correct answer: act on nothing.
PRESET_PLAN_GUARDRAIL = {
    "candidates": [
        {
            "candidate_id": "restock:005-00",
            "operation_type": "restock_inventory",
            "facts": {
                "committed": 2,
                "daily_velocity": 14.7143,
                "data_complete": True,
                "eligible": True,
                "inbound": 3,
                "lead_time_days": 18,
                "minimum_order_quantity": 24,
                "net_available": 1,
                "on_hand": 0,
                "pack_size": 5,
                "reorder_point": 271,
                "reorder_quantity": 420,
                "safety_stock": 6,
                "sales_lookback_days": 14,
                "sales_units": 206,
                "supplier_configured": True,
                "target_coverage_days": 28,
                "target_stock": 418,
                "validated_quantity": 420,
            },
        }
    ],
    "enabled_operations": ["flag_issue"],
}

# --- model ---------------------------------------------------------------

print(f"Loading {BASE_MODEL} + adapter {ADAPTER} …")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
# Qwen3.5-9B is Qwen3_5ForConditionalGeneration (text_config + vision_config).
# It MUST be loaded through the conditional-generation class: that nests the
# decoder under model.language_model.*, which is the module path the adapter
# checkpoint was trained against. AutoModelForCausalLM flattens it to
# model.layers.*, PEFT then matches zero keys, leaves lora_B at its zero init,
# and the adapter silently becomes an identity function.
model = AutoModelForImageTextToText.from_pretrained(BASE_MODEL, dtype=torch.bfloat16)
# torch_device="cpu" is REQUIRED under ZeroGPU. torch.cuda.is_available() is
# patched to True at startup, so PEFT would otherwise infer device="cuda" and
# hand it straight to safetensors — which fails, because no GPU is actually
# allocated outside a @spaces.GPU function.
model = PeftModel.from_pretrained(model, ADAPTER, torch_device="cpu")
model.eval()

# Guard: lora_B is zero-initialised by construction, so a nonzero norm is proof
# that trained weights actually bound to the module tree. Without this check a
# key mismatch is invisible and every generation is silently base-model-only.
_b = [float(p.detach().float().norm()) for n, p in model.named_parameters() if "lora_B" in n]
ADAPTER_BOUND = sum(1 for x in _b if x > 0)
print(f"adapter check: {ADAPTER_BOUND}/{len(_b)} lora_B tensors nonzero")
if not ADAPTER_BOUND:
    print("FATAL: adapter did NOT bind (key mismatch) — outputs would be base model only.")

# ZeroGPU intercepts this .to("cuda") and materialises it when a GPU is granted.
if ZERO_GPU or torch.cuda.is_available():
    model = model.to("cuda")
print(f"Model ready on {model.device}.")


@GPU
def _generate(task: str, payload: str, max_new_tokens: int, thinking: bool = False) -> str:
    messages = [
        {"role": "system", "content": SYSTEM[task]},
        {"role": "user", "content": payload},
    ]
    # Training targets contain zero <think> blocks, so enable_thinking=False matches
    # the training distribution. But the base Qwen3.5 template opens a <think> block
    # by default, and the fine-tune may have been run that way — hence the toggle.
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=thinking,
    )
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,  # structured extraction — greedy, not creative
            pad_token_id=tokenizer.convert_tokens_to_ids("<|endoftext|>"),
            eos_token_id=tokenizer.convert_tokens_to_ids("<|im_end|>"),
        )
    completion = out[0][inputs["input_ids"].shape[-1] :]
    return tokenizer.decode(completion, skip_special_tokens=True).strip()


@GPU
def _generate_base(task: str, payload: str, max_new_tokens: int) -> str:
    """Same prompt with the LoRA adapter disabled — the control arm."""
    text = tokenizer.apply_chat_template(
        [{"role": "system", "content": SYSTEM[task]}, {"role": "user", "content": payload}],
        tokenize=False, add_generation_prompt=True, enable_thinking=False,
    )
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    with torch.no_grad(), model.disable_adapter():
        out = model.generate(
            **inputs, max_new_tokens=max_new_tokens, do_sample=False,
            pad_token_id=tokenizer.convert_tokens_to_ids("<|endoftext|>"),
            eos_token_id=tokenizer.convert_tokens_to_ids("<|im_end|>"),
        )
    completion = out[0][inputs["input_ids"].shape[-1] :]
    return tokenizer.decode(completion, skip_special_tokens=True).strip()


# --- guardrail validation -------------------------------------------------


def _field(item, key: str):
    """Read a field from a response item.

    Training data emits objects — {"operation_type": ..., "reason": ...} — but the
    model sometimes emits a bare string instead. Accept both so a schema deviation
    is reported as a finding rather than crashing the validator.
    """
    if isinstance(item, dict):
        return item.get(key)
    if isinstance(item, str):
        return item if key == "operation_type" else None
    return None


def _validate(task: str, payload_obj: dict, raw: str) -> str:
    """Check the response against the constraints the model was trained on."""
    lines = []
    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError as e:
        return f"❌ **Invalid JSON** — {e}"
    if not isinstance(parsed, dict):
        return f"❌ **Expected a JSON object, got `{type(parsed).__name__}`**"
    lines.append("✅ **Valid JSON**")

    key = "recommended_operations" if task == "capability_advice" else "operations"
    items = parsed.get(key)
    if items is None:
        return "\n\n".join(lines + [f"❌ **Missing `{key}` field**"])
    if not isinstance(items, list):
        return "\n\n".join(lines + [f"❌ **`{key}` is not a list** (got `{type(items).__name__}`)"])

    # Flag the shape deviation explicitly — it is real signal about the model.
    if any(isinstance(i, str) for i in items):
        lines.append(
            f"⚠️ **Schema deviation:** `{key}` contains bare strings; "
            "training data used objects with `operation_type`."
        )

    if task == "capability_advice":
        allowed = set(payload_obj.get("available_operations", []))
        bad = [t for t in (_field(i, "operation_type") for i in items) if t not in allowed]
        lines.append(
            f"❌ **Unregistered operations recommended:** {bad}"
            if bad
            else f"✅ **All {len(items)} recommendation(s) are registered operations**"
        )
    else:
        enabled = set(payload_obj.get("enabled_operations", []))
        valid_ids = {c.get("candidate_id") for c in payload_obj.get("candidates", []) if isinstance(c, dict)}

        not_enabled = [t for t in (_field(i, "operation_type") for i in items) if t not in enabled]
        lines.append(
            f"❌ **Selected non-enabled operations:** {not_enabled}"
            if not_enabled
            else "✅ **Every selected operation is enabled**"
        )

        ids = [_field(i, "candidate_id") for i in items]
        invented = [c for c in ids if c not in valid_ids]
        lines.append(
            f"❌ **Invented candidate IDs:** {invented}"
            if invented
            else "✅ **No invented candidate IDs**"
        )

        if not items:
            lines.append("ℹ️ Model declined to act — correct when no candidate is enabled.")

    return "\n\n".join(lines)


def run(task: str, payload: str, max_new_tokens: int):
    try:
        payload_obj = json.loads(payload)
    except json.JSONDecodeError as e:
        return "(not run — input is not valid JSON)", f"❌ **Input JSON error** — {e}"

    compact = json.dumps(payload_obj, separators=(",", ":"))  # match training format
    try:
        raw = _generate(task, compact, int(max_new_tokens))
    except Exception as e:  # GPU allocation, OOM, decoding — surface, never crash
        return "(generation failed)", f"❌ **Generation error** — `{type(e).__name__}`: {e}"

    try:  # pretty-print if parseable
        pretty = json.dumps(json.loads(raw), indent=2)
    except json.JSONDecodeError:
        pretty = raw  # show exactly what the model emitted

    try:
        verdict = _validate(task, payload_obj, raw)
    except Exception as e:  # a validator bug must never hide the model's output
        verdict = f"⚠️ **Validator error** — `{type(e).__name__}`: {e}"
    return pretty, verdict


# --- UI ------------------------------------------------------------------

with gr.Blocks(title="Commerce Operations Agent") as demo:
    gr.Markdown(
        "# 🛒 Commerce Operations Agent\n"
        f"LoRA adapter [`{ADAPTER}`](https://huggingface.co/{ADAPTER}) on `{BASE_MODEL}`.\n\n"
        "A **structured-JSON transducer** for e-commerce operations — not a chat model. "
        "The right-hand panel validates each response against the guardrails the model "
        "was trained on."
    )

    with gr.Tab("Capability advice"):
        gr.Markdown(
            "*Given a store summary, which operations are worth enabling?* "
            "The model may only recommend operations listed in `available_operations`."
        )
        with gr.Row():
            with gr.Column():
                cap_in = gr.Code(
                    value=json.dumps(PRESET_CAPABILITY, indent=2),
                    language="json",
                    label="Input",
                    lines=16,
                )
                cap_tokens = gr.Slider(64, 1024, value=512, step=64, label="Max new tokens")
                cap_btn = gr.Button("Run", variant="primary")
            with gr.Column():
                cap_out = gr.Code(language="json", label="Model output", lines=16)
                cap_valid = gr.Markdown(label="Guardrail check")
        cap_btn.click(
            lambda p, t: run("capability_advice", p, t),
            [cap_in, cap_tokens],
            [cap_out, cap_valid],
        )

    with gr.Tab("Operation plan"):
        gr.Markdown(
            "*Given deterministic candidates, which does the agent act on?* "
            "The model may only select candidates whose `operation_type` is in "
            "`enabled_operations`. **Try the guardrail preset** — the candidate is urgent, "
            "but its operation is not enabled, so the correct answer is to act on nothing."
        )
        with gr.Row():
            with gr.Column():
                plan_in = gr.Code(
                    value=json.dumps(PRESET_PLAN_ACT, indent=2),
                    language="json",
                    label="Input",
                    lines=22,
                )
                with gr.Row():
                    act_btn = gr.Button("Load: should act")
                    guard_btn = gr.Button("Load: guardrail (should decline)")
                plan_tokens = gr.Slider(64, 1024, value=512, step=64, label="Max new tokens")
                plan_btn = gr.Button("Run", variant="primary")
            with gr.Column():
                plan_out = gr.Code(language="json", label="Model output", lines=22)
                plan_valid = gr.Markdown(label="Guardrail check")

        act_btn.click(lambda: json.dumps(PRESET_PLAN_ACT, indent=2), None, plan_in)
        guard_btn.click(lambda: json.dumps(PRESET_PLAN_GUARDRAIL, indent=2), None, plan_in)
        plan_btn.click(
            lambda p, t: run("operation_plan", p, t),
            [plan_in, plan_tokens],
            [plan_out, plan_valid],
        )

    gr.Markdown(
        "---\n"
        "Decoding is greedy — this is extraction, not generation. Thinking mode defaults to "
        "off (training targets contain no `<think>` blocks), but the `generate` API endpoint "
        "accepts a `thinking` flag so both modes can be compared."
    )

    # Stable endpoint for batch evaluation: returns the model's unmodified output.
    def api_generate(
        task: str, payload: str, thinking: bool = False, max_new_tokens: int = 512
    ) -> str:
        compact = json.dumps(json.loads(payload), separators=(",", ":"))
        return _generate(task, compact, int(max_new_tokens), thinking)

    gr.api(api_generate, api_name="generate")

    def api_diagnose(payload: str, task: str = "operation_plan") -> dict:
        """Prove whether the LoRA adapter is actually affecting generation.

        Runs the identical prompt with the adapter active and with it disabled.
        Identical outputs would mean the fine-tune is not in the loop at all and
        every evaluation number is measuring the base model.
        """
        compact = json.dumps(json.loads(payload), separators=(",", ":"))
        prompt = tokenizer.apply_chat_template(
            [{"role": "system", "content": SYSTEM[task]},
             {"role": "user", "content": compact}],
            tokenize=False, add_generation_prompt=True, enable_thinking=False,
        )
        with_adapter = _generate(task, compact, 256, False)
        without_adapter = _generate_base(task, compact, 256)

        cfg = {}
        try:
            pc = list(model.peft_config.values())[0]
            cfg = {"r": pc.r, "lora_alpha": pc.lora_alpha,
                   "target_modules": sorted(pc.target_modules) if pc.target_modules else None,
                   "task_type": str(pc.task_type)}
        except Exception as e:
            cfg = {"error": str(e)}

        lora_layers = [n for n, _ in model.named_modules() if "lora_A" in n]
        return {
            "exact_prompt_sent": prompt,
            "peft_config": cfg,
            "active_adapters": getattr(model, "active_adapters", None) if isinstance(
                getattr(model, "active_adapters", None), list) else str(getattr(model, "active_adapter", None)),
            "lora_layer_count": len(lora_layers),
            "adapter_bound_tensors": ADAPTER_BOUND,
            "sample_lora_layers": lora_layers[:3],
            "model_device": str(model.device),
            "output_WITH_adapter": with_adapter,
            "output_WITHOUT_adapter": without_adapter,
            "outputs_identical": with_adapter.strip() == without_adapter.strip(),
        }

    gr.api(api_diagnose, api_name="diagnose")

if __name__ == "__main__":
    demo.launch()