pptx_tool.py

"""
pptx_tool.py — AIToolkit tool-call adapter for the ppt-llm presentation engine.

Exposes LLM-callable tools backed by the vendored ``pptx_engine`` package:

  • generate_presentation — topic → beautiful, editable .pptx (+ inline preview)
  • revise_presentation   — natural-language edits to a previously generated deck
  • list_themes           — discover available themes

Design (see pptx_engine): the chat model only emits a validated JSON deck-spec;
a deterministic python-pptx renderer turns it into the file. We route the
engine's LLM calls through AIToolkit's provider layer (get_client) so the user's
configured provider / model / key and EU-mode are honoured.

Wired into tool_executor.py via PPTX_TOOL_SCHEMAS (→ TOOL_REGISTRY); the
stateful handlers are built per-request in chat_handlers.py through
make_pptx_tool_handlers(api_keys, user_state, provider, model, key).
"""

from __future__ import annotations

import base64
import logging
import os
import re
import shutil
import time
import uuid

logger = logging.getLogger(__name__)

# Gradio serves files from here (see app.py allowed_paths); per-user storage is
# NOT allow-listed, so we publish a download copy into this dir.
_DOWNLOAD_DIR = os.environ.get("PPTLLM_DOWNLOAD_DIR", "/tmp/gradio_downloads")  # nosec B108

try:
    from pptx_engine.assets import (
        FunctionProvider,
        NullProvider,
        get_provider,
        resolve_deck_images,
    )
    from pptx_engine.planner import Brief, generate_presentation, revise_deck
    from pptx_engine.preview import generate_previews, soffice_available
    from pptx_engine.render import render_deck
    from pptx_engine.spec import Deck, validate_budgets  # noqa: F401
    from pptx_engine.themes import THEMES
    from pptx_engine.tools import TOOL_SCHEMAS as _ENGINE_SCHEMAS

    HAS_PPTX_ENGINE = True
except ImportError as e:  # pragma: no cover
    logger.warning(f"[pptx_tool] pptx_engine not importable: {e}")
    HAS_PPTX_ENGINE = False

# Default vision model for an (optional) QA pass, only used with Mistral.
_VISION_MODEL = os.environ.get("PPTLLM_VISION_MODEL", "pixtral-12b-2409")


def _storage_root() -> str:
    try:
        from config import STORAGE_MOUNT_POINT

        return STORAGE_MOUNT_POINT
    except Exception:
        return os.path.join(os.path.dirname(__file__), "storage")


def _user_dir(user_state: dict | None) -> str:
    uid = (user_state or {}).get("id", "anon")
    d = os.path.join(_storage_root(), f"user_{uid}", "presentations")
    os.makedirs(d, exist_ok=True)
    return d


# ── LLM adapter: satisfy the engine's chat_json/chat interface via get_client ──
class _ToolkitLLM:
    """Adapts AIToolkit's OpenAI-compatible client to the engine's LLM API."""

    def __init__(self, client, model: str, vision_model: str | None = None):
        self.client = client
        self.model = model
        self.vision_model = vision_model or model

    def _has_image(self, messages) -> bool:
        for m in messages:
            c = m.get("content")
            if isinstance(c, list) and any(
                isinstance(p, dict) and p.get("type") == "image_url" for p in c
            ):
                return True
        return False

    def chat(self, messages, *, json_mode: bool = True, max_tokens: int = 4000) -> str:
        kwargs = {
            "model": self.vision_model if self._has_image(messages) else self.model,
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": 0.4,
        }
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}
        resp = self.client.chat.completions.create(**kwargs)
        return resp.choices[0].message.content

    def chat_json(self, messages, *, max_tokens: int = 4000) -> dict:
        from pptx_engine.llm import _parse_json

        return _parse_json(self.chat(messages, json_mode=True, max_tokens=max_tokens))


def _deck_summary(deck: Deck, deck_id: str, pptx_path: str) -> str:
    hist: dict[str, int] = {}
    for s in deck.slides:
        hist[s.layout.value] = hist.get(s.layout.value, 0) + 1
    layouts = ", ".join(f"{k}×{v}" for k, v in hist.items())
    return (
        f"Generated '{deck.title}' — {len(deck.slides)} slides ({layouts}), "
        f"theme '{deck.theme}'. deck_id={deck_id}. Saved to {pptx_path}."
    )


def _publish_download(pptx_path: str, deck_id: str, title: str) -> str:
    """Copy the deck into the Gradio-served downloads dir and return a markdown
    download link (the [name](/file=...) form the app already uses)."""
    try:
        safe = re.sub(r"[^\w\-]+", "_", title).strip("_")[:60] or "presentation"
        dest_dir = os.path.join(_DOWNLOAD_DIR, deck_id)
        os.makedirs(dest_dir, exist_ok=True)
        dest = os.path.join(dest_dir, f"{safe}.pptx")
        shutil.copyfile(pptx_path, dest)
        return f"📎 [{safe}.pptx](/file={dest})"
    except Exception as e:  # download link is non-fatal
        logger.debug(f"[pptx_tool] publish download failed: {e}")
        return f"📎 saved to <code>{pptx_path}</code>"


def _inline_preview(pptx_path: str, out_dir: str) -> str:
    """Best-effort first-slide PNG embedded as inline HTML for the chat bubble."""
    if not soffice_available():
        return ""
    try:
        pngs = generate_previews(pptx_path, os.path.join(out_dir, "_preview"), dpi=96, max_slides=1)
        if not pngs:
            return ""
        with open(pngs[0], "rb") as f:
            b64 = base64.b64encode(f.read()).decode("ascii")
        return (
            f'<img src="data:image/png;base64,{b64}" '
            f'style="max-width:100%;border-radius:8px;margin-top:8px;" />'
        )
    except Exception as e:  # previews are non-fatal
        logger.debug(f"[pptx_tool] preview failed: {e}")
        return ""


# ── handler factory (stateful: needs provider/model/key + user) ───────────────
def make_pptx_tool_handlers(
    api_keys: dict, user_state: dict | None, provider: str, model: str, key: str | None
):
    """Return {tool_name: handler(args)->ToolResult}. Mirrors _make_image_gen_handler."""
    if not HAS_PPTX_ENGINE:
        return {}

    def _llm() -> _ToolkitLLM:
        from provider_utils import get_client

        client = get_client(provider, key)
        vision = _VISION_MODEL if str(provider).lower().startswith("mistral") else model
        return _ToolkitLLM(client, model, vision)

    def _result(text: str, ui: str | None = None, error: bool = False):
        from tools import ToolResult  # lazy to avoid import cycle

        return ToolResult(text=text, ui=ui, error=error)

    def _image_provider(image_mode: str):
        """stock = Pexels/Picsum; generate = AI imagery via the app's EU FLUX
        (BFL, else Nebius) through run_image_gen; none = placeholder."""
        mode = (image_mode or "stock").lower()
        if mode == "none":
            return NullProvider()
        if mode in ("generate", "flux"):
            from image_gen_utils import run_image_gen

            bfl = api_keys.get("BFL", "")
            neb = api_keys.get("NEBIUS", "")
            if bfl:
                gp, gm, gkey, steps = "BFL", "flux-2-klein-4b", bfl, 4
            elif neb:
                gp, gm, gkey, steps = "Nebius", "black-forest-labs/flux-schnell", neb, 4
            else:
                return get_provider()  # no gen key → fall back to stock

            def _gen(query, w, h):
                path, _status = run_image_gen(query, gp, gm, w, h, steps, gkey, user_state)
                return path

            return FunctionProvider(_gen, name=f"gen:{gp}")
        return get_provider()  # stock

    def generate(args: dict):
        topic = (args.get("topic") or "").strip()
        if not topic:
            return _result("Error: 'topic' is required.", error=True)
        theme = args.get("theme") or "midnight"
        if theme not in THEMES:
            theme = "midnight"
        brief = Brief(
            topic=topic,
            audience=args.get("audience") or "a general professional audience",
            num_slides=int(args.get("num_slides") or 8),
            theme=theme,
            language=args.get("language") or "en",
            tone=args.get("tone") or "clear, confident, concrete",
        )
        image_mode = (args.get("image_mode") or "stock").lower()
        try:
            res = generate_presentation(brief, _llm())
            if image_mode != "none":
                resolve_deck_images(res.deck, provider=_image_provider(image_mode))
            d = _user_dir(user_state)
            deck_id = f"{int(time.time())}-{uuid.uuid4().hex[:8]}"
            res.deck.model_dump_json()  # validate serialisable
            with open(os.path.join(d, f"{deck_id}.json"), "w", encoding="utf-8") as f:
                f.write(res.deck.model_dump_json(indent=2))
            pptx_path = os.path.join(d, f"{deck_id}.pptx")
            render_deck(res.deck, pptx_path)
        except Exception as e:
            logger.exception("[pptx_tool] generate failed")
            return _result(f"Presentation generation failed: {e}", error=True)
        text = _deck_summary(res.deck, deck_id, pptx_path)
        link = _publish_download(pptx_path, deck_id, res.deck.title)
        ui = f"✅ {text}<br/>{link}" + _inline_preview(pptx_path, d)
        return _result(text, ui=ui)

    def revise(args: dict):
        deck_id = (args.get("deck_id") or "").strip()
        instructions = (args.get("instructions") or "").strip()
        if not deck_id or not instructions:
            return _result("Error: 'deck_id' and 'instructions' are required.", error=True)
        # Sanitize deck_id to prevent path traversal
        import re as _re

        if not _re.match(r"^[\w-]+$", deck_id):
            return _result("Error: invalid deck_id format.", error=True)
        d = _user_dir(user_state)
        spec_path = os.path.join(d, f"{deck_id}.json")
        if not os.path.exists(spec_path):
            return _result(f"Error: unknown deck_id '{deck_id}'.", error=True)
        try:
            with open(spec_path, encoding="utf-8") as f:
                deck = Deck.model_validate_json(f.read())
            revised = revise_deck(_llm(), deck, instructions)
            image_mode = (args.get("image_mode") or "stock").lower()
            if image_mode != "none":
                resolve_deck_images(revised, provider=_image_provider(image_mode))
            with open(spec_path, "w", encoding="utf-8") as f:
                f.write(revised.model_dump_json(indent=2))
            pptx_path = os.path.join(d, f"{deck_id}.pptx")
            render_deck(revised, pptx_path)
        except Exception as e:
            logger.exception("[pptx_tool] revise failed")
            return _result(f"Revision failed: {e}", error=True)
        text = _deck_summary(revised, deck_id, pptx_path)
        link = _publish_download(pptx_path, deck_id, revised.title)
        ui = f"✅ Revised. {text}<br/>{link}" + _inline_preview(pptx_path, d)
        return _result(text, ui=ui)

    def list_themes(args: dict):
        names = ", ".join(THEMES.keys())
        return _result(f"Available themes: {names}.")

    def evaluate(args: dict):
        deck_id = (args.get("deck_id") or "").strip()
        if not deck_id:
            return _result("Error: 'deck_id' is required.", error=True)
        pptx_path = os.path.join(_user_dir(user_state), f"{deck_id}.pptx")
        if not os.path.exists(pptx_path):
            return _result(f"Error: unknown deck_id '{deck_id}'.", error=True)
        if not soffice_available():
            return _result(
                "QA needs LibreOffice (preview rendering) which is not installed on the server.",
                error=True,
            )
        mistral_key = api_keys.get("MISTRAL", "")
        if not mistral_key:
            return _result("QA needs a Mistral key for the vision model.", error=True)
        try:
            from pptx_engine.qa import evaluate_deck, weak_slides
            from provider_utils import get_client

            client = get_client("Mistral", mistral_key)
            vis = _ToolkitLLM(client, _VISION_MODEL, _VISION_MODEL)
            report = evaluate_deck(pptx_path, llm=vis)
        except Exception as e:
            logger.exception("[pptx_tool] evaluate failed")
            return _result(f"QA evaluation failed: {e}", error=True)
        s = report.summary()
        weak = [w.index + 1 for w in weak_slides(report)]
        text = (
            f"QA scores for deck {deck_id} (1-5): overall {s['overall']}, "
            f"content {s['content']}, design {s['design']}, "
            f"coherence {s['coherence']}. "
            + (f"Weak slides: {weak}." if weak else "No weak slides.")
        )
        rows = "".join(
            f"<li>Slide {sl['index'] + 1}: content {sl['content']}, "
            f"design {sl['design']} — {sl['notes']}</li>"
            for sl in s["slides"]
        )
        ui = (
            f"📊 <b>QA — overall {s['overall']}/5</b> "
            f"(content {s['content']}, design {s['design']}, "
            f"coherence {s['coherence']})<br/><ul>{rows}</ul>"
            f"<i>{s['coherence_notes']}</i>"
        )
        return _result(text, ui=ui)

    return {
        "generate_presentation": generate,
        "revise_presentation": revise,
        "evaluate_presentation": evaluate,
        "list_themes": list_themes,
    }


# ── schema exports for tool_executor.py ───────────────────────────────────────
if HAS_PPTX_ENGINE:
    PPTX_TOOL_SCHEMAS: dict[str, dict] = {
        s["function"]["name"]: s
        for s in _ENGINE_SCHEMAS
        if s["function"]["name"]
        in {"generate_presentation", "revise_presentation", "evaluate_presentation", "list_themes"}
    }
else:  # pragma: no cover
    PPTX_TOOL_SCHEMAS = {}

PPTX_TOOL_NAMES = list(PPTX_TOOL_SCHEMAS.keys())