hkuds--deeptutor/deeptutor/tools/media_gen_tool.py

"""Image- and video-generation chat tools.

Thin BaseTool front-ends over the ``imagegen`` / ``videogen`` services. Each
call generates media via the active catalog model, writes the bytes into the
turn's public workspace, and returns the files as artifacts — reusing the exact
same ``collect_public_artifacts`` convention as the exec / code_execution tools,
so generated images/videos surface in chat as download cards (and the model can
linkify them by writing the filename verbatim).

Mounting & gating (set by the chat pipeline, not here):

* User-toggleable in /settings/tools, so per-user ``enabled_tools`` grants apply.
* Only mounted for a turn when an active model is configured for the service.
* ``_workspace_dir`` is injected server-side by ``_augment_tool_kwargs``; the
  LLM only supplies ``prompt`` (+ optional knobs).
"""

from __future__ import annotations

import logging
from pathlib import Path
import re
from typing import TYPE_CHECKING, Any
import uuid

from deeptutor.core.tool_protocol import BaseTool, ToolDefinition, ToolParameter, ToolResult

if TYPE_CHECKING:
    from deeptutor.services.sandbox.artifacts import SandboxArtifact

logger = logging.getLogger(__name__)

_EXT_BY_CONTENT_TYPE = {
    "image/png": "png",
    "image/jpeg": "jpg",
    "image/webp": "webp",
    "image/gif": "gif",
    "video/mp4": "mp4",
    "video/webm": "webm",
    "video/quicktime": "mov",
}


def _slug(prompt: str, fallback: str) -> str:
    """Short ascii filename stem from a prompt; ``fallback`` when none survives."""
    words = re.findall(r"[A-Za-z0-9]+", prompt.lower())
    stem = "_".join(words[:6])[:40]
    return stem or fallback


def _ext(content_type: str, default: str) -> str:
    return _EXT_BY_CONTENT_TYPE.get((content_type or "").split(";")[0].strip(), default)


def _run_dir(injected: str | None, kind: str) -> Path:
    """A fresh per-call dir under the public outputs root.

    Uses the pipeline-injected workspace when present; falls back to the same
    public chat media workspace for direct/tool-test calls. Per-call subdir
    keeps ``collect_public_artifacts`` scoped to this call's files only.
    """
    if injected:
        base = Path(injected)
    else:
        from deeptutor.services.path_service import get_path_service

        base = get_path_service().get_task_workspace("chat", "media_gen") / "media"
    run = base / f"{kind}_{uuid.uuid4().hex[:12]}"
    run.mkdir(parents=True, exist_ok=True)
    return run


def _write_media(
    run_dir: Path,
    media: list[tuple[bytes, str]],
    *,
    stem: str,
    default_ext: str,
) -> list[SandboxArtifact]:
    """Write generated bytes to ``run_dir`` and return their public artifacts."""
    from deeptutor.services.sandbox.artifacts import collect_public_artifacts

    multiple = len(media) > 1
    for index, (data, content_type) in enumerate(media, start=1):
        suffix = f"_{index}" if multiple else ""
        filename = f"{stem}{suffix}.{_ext(content_type, default_ext)}"
        (run_dir / filename).write_bytes(data)
    return collect_public_artifacts(str(run_dir))


def _artifact_result(
    artifacts: list[SandboxArtifact], *, empty_message: str, **meta: Any
) -> ToolResult:
    from deeptutor.services.sandbox.artifacts import render_artifacts_for_tool

    if not artifacts:
        return ToolResult(content=empty_message, success=False)
    rows = [artifact.to_dict() for artifact in artifacts]
    return ToolResult(
        content=render_artifacts_for_tool(artifacts),
        sources=[
            {
                "type": "artifact",
                "filename": row["filename"],
                "url": row["url"],
                "path": row["path"],
                "mime_type": row["mime_type"],
                "size_bytes": row["size_bytes"],
            }
            for row in rows
        ],
        metadata={"artifacts": rows, **meta},
    )


class ImagegenTool(BaseTool):
    """Generate images from a text prompt via the configured imagegen model."""

    def get_definition(self) -> ToolDefinition:
        return ToolDefinition(
            name="imagegen",
            description=(
                "Generate one or more images from a text description using the "
                "configured image-generation model. Write a vivid, self-contained "
                "`prompt` describing the subject, style, and composition. Generated "
                "images are saved and shown to the user automatically as cards — "
                "after calling, refer to each image by its exact filename in your reply."
            ),
            parameters=[
                ToolParameter(
                    name="prompt",
                    type="string",
                    description="A detailed description of the image to generate.",
                ),
                ToolParameter(
                    name="size",
                    type="string",
                    description="Optional WxH like '1024x1024' or '1792x1024'. Omit for the default.",
                    required=False,
                ),
                ToolParameter(
                    name="n",
                    type="integer",
                    description="How many images to generate (1-4). Default 1.",
                    required=False,
                    default=1,
                ),
            ],
        )

    async def execute(self, **kwargs: Any) -> ToolResult:
        from deeptutor.services.imagegen import GenerationProviderError, generate_image

        prompt = str(kwargs.get("prompt") or "").strip()
        if not prompt:
            return ToolResult(content="imagegen requires a non-empty 'prompt'.", success=False)
        size = str(kwargs.get("size") or "").strip() or None
        try:
            count = int(kwargs.get("n") or 1)
        except (TypeError, ValueError):
            count = 1
        count = max(1, min(count, 4))

        try:
            images = await generate_image(prompt, size=size, n=count)
        except ValueError as exc:  # not configured
            return ToolResult(content=str(exc), success=False)
        except GenerationProviderError as exc:
            logger.warning("imagegen failed: %s", exc)
            return ToolResult(content=f"Image generation failed: {exc}", success=False)

        run_dir = _run_dir(kwargs.get("_workspace_dir"), "imagegen")
        artifacts = _write_media(run_dir, images, stem=_slug(prompt, "image"), default_ext="png")
        return _artifact_result(
            artifacts,
            empty_message="Image generation produced no saved files.",
            prompt=prompt,
            count=len(images),
            kind="image",
        )


class VideogenTool(BaseTool):
    """Generate a video from a text prompt via the configured videogen model."""

    def get_definition(self) -> ToolDefinition:
        return ToolDefinition(
            name="videogen",
            description=(
                "Generate a short video from a text description using the configured "
                "video-generation model. Write a vivid, self-contained `prompt` "
                "describing the scene, motion, and style. Rendering is slow (it may "
                "take a minute or more). The video is saved and shown to the user "
                "automatically — after calling, refer to it by its exact filename."
            ),
            parameters=[
                ToolParameter(
                    name="prompt",
                    type="string",
                    description="A detailed description of the video to generate.",
                ),
                ToolParameter(
                    name="aspect_ratio",
                    type="string",
                    description="Optional aspect ratio like '16:9' or '9:16'. Omit for the default.",
                    required=False,
                ),
                ToolParameter(
                    name="duration",
                    type="string",
                    description="Optional duration in seconds (e.g. '5'). Omit for the default.",
                    required=False,
                ),
            ],
        )

    async def execute(self, **kwargs: Any) -> ToolResult:
        from deeptutor.services.videogen import GenerationProviderError, generate_video

        prompt = str(kwargs.get("prompt") or "").strip()
        if not prompt:
            return ToolResult(content="videogen requires a non-empty 'prompt'.", success=False)
        aspect_ratio = str(kwargs.get("aspect_ratio") or "").strip() or None
        duration = str(kwargs.get("duration") or "").strip() or None

        event_sink = kwargs.get("event_sink")

        async def _progress(message: str) -> None:
            if event_sink is None:
                return
            try:
                await event_sink("tool_log", message)
            except Exception:  # progress is best-effort; never fail the render
                logger.debug("videogen progress emit failed", exc_info=True)

        try:
            video, content_type = await generate_video(
                prompt,
                aspect_ratio=aspect_ratio,
                duration=duration,
                progress=_progress,
            )
        except ValueError as exc:  # not configured
            return ToolResult(content=str(exc), success=False)
        except GenerationProviderError as exc:
            logger.warning("videogen failed: %s", exc)
            return ToolResult(content=f"Video generation failed: {exc}", success=False)

        run_dir = _run_dir(kwargs.get("_workspace_dir"), "videogen")
        artifacts = _write_media(
            run_dir, [(video, content_type)], stem=_slug(prompt, "video"), default_ext="mp4"
        )
        return _artifact_result(
            artifacts,
            empty_message="Video generation produced no saved file.",
            prompt=prompt,
            kind="video",
        )


__all__ = ["ImagegenTool", "VideogenTool"]