"""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"]