api/uxmaster.py

"""
UX MASTER endpoint — POST /api/uxmaster.

A multimodal design-critique agent: the user uploads a SCREENSHOT of their running app;
UX MASTER critiques the ACTUAL rendered UI and returns the complete improved file set via a
single vision LLM call (Gemini 2.5 Flash). Mirrors /api/enhance (premium gate, parse → reply,
normalize → upsert) but adds a `needs_verification` flag the agent uses to ask for a fresh
screenshot.
"""
import logging

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel

from agent.llm import VISION_CHAIN, chat_vision
from agent.premium_scaffold import is_premium_shaped, preserve_premium_scaffold
from agent.prompts import CODER_SYSTEM_PROMPT, uxmaster_prompt
from agent.scaffold import (detect_framework, normalize_files,
                           parse_uxmaster_response)
from agent.llm import strip_code_fences
from api.enhance import _upsert_project  # ponytail: reuse the same owned-project upsert
from auth.dependencies import get_current_user
from db.models import User

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/uxmaster", tags=["uxmaster"])


class UXMasterRequest(BaseModel):
    project_id: str | None = None
    files: dict[str, str]
    image: str  # base64 data URL of the running-app screenshot
    instruction: str | None = None
    session_id: str | None = None  # ponytail: kept for backward compat, ignored for ownership
    premium: bool = False  # honored ONLY for SUPER_ADMIN (server-side gate below)


class UXMasterResponse(BaseModel):
    project_id: str
    files: dict[str, str]
    reply: str  # first-person message from UX MASTER (what it saw + changed)
    needs_verification: bool  # True → UX MASTER wants a fresh screenshot to verify


@router.post("", response_model=UXMasterResponse)
def uxmaster(
    req: UXMasterRequest,
    current_user: User = Depends(get_current_user),
) -> UXMasterResponse:
    # ponytail: validate at the boundary — the screenshot is the whole point of this agent.
    if not req.image or not req.image.strip().startswith("data:"):
        raise HTTPException(status_code=400, detail="image must be a base64 data URL")
    if not req.files:
        raise HTTPException(status_code=400, detail="files must not be empty")

    instruction = (req.instruction or "").strip()
    # ponytail: never trust the client flag — premium is honored only for SUPER_ADMIN (mirrors enhance).
    premium = bool(req.premium) and getattr(current_user, "role", "") == "SUPER_ADMIN"
    raw = chat_vision(
        uxmaster_prompt(req.files, instruction, image_provided=True),
        req.image,
        chain=VISION_CHAIN,  # gemini-only — the multimodal model in the pool
        system=CODER_SYSTEM_PROMPT,
        max_tokens=20000 if premium else 8192, agent="uxmaster",
    )
    # ponytail: ONE call emits the conversational reply, the verify boolean, and the files.
    reply, needs_verification, parsed = parse_uxmaster_response(raw, strip_code_fences)
    reply = reply or "Reviewed your screenshot and refined the UI."
    updated_files = normalize_files(parsed, detect_framework(req.files))
    # ponytail: same premium-safety as enhance — never drop/clobber the baked /ui kit.
    if is_premium_shaped(req.files):
        updated_files = preserve_premium_scaffold(req.files, updated_files)

    label = instruction or "screenshot review"
    project_id = _upsert_project(req.project_id, current_user.id, f"UX Master: {label}", updated_files, source="uxmaster")
    return UXMasterResponse(
        project_id=project_id,
        files=updated_files,
        reply=reply,
        needs_verification=needs_verification,
    )