add slide translation

mcp/src/aileen3_mcp/media_tools.py

@@ -16,6 +16,7 @@ from typing import Any, Callable, Dict, Optional
 
 import ffmpeg
 from fastmcp import Context, FastMCP
+from mcp.types import ImageContent
 from contextlib import redirect_stdout, redirect_stderr, contextmanager
 import io
 from PIL import Image
@@ -38,7 +39,6 @@ DEBUG = os.environ.get("AILEEN3_DEBUG", "").lower() in {"1", "true", "yes", "on"
 DEBUG_DIR = Path(tempfile.gettempdir()) / "aileen3-debug"
 if DEBUG:
     DEBUG_DIR.mkdir(parents=True, exist_ok=True)
-
 def _write_debug(reference: str, suffix: str, data: Any) -> None:
     if not DEBUG:
         return
@@ -602,6 +602,108 @@ def _gemini_analyze_audio(client, audio_path: Path, slides: list[dict], priors:
     }
 
 
+def _language_slug(value: str) -> str:
+    value = (value or "").strip().lower()
+    value = re.sub(r"[^a-z0-9]+", "-", value)
+    value = value.strip("-")
+    return value or "translation"
+
+
+def _slide_image_bytes(reference: str, slide: dict) -> bytes | None:
+    data_uri = slide.get("image_data_uri")
+    if isinstance(data_uri, str) and data_uri.startswith("data:"):
+        try:
+            _, payload = data_uri.split(",", 1)
+            return base64.b64decode(payload)
+        except Exception:
+            pass
+
+    idx = slide.get("index")
+    if idx is not None:
+        try:
+            idx_int = int(idx)
+        except Exception:
+            idx_int = None
+        if idx_int is not None:
+            path = SLIDE_CACHE / reference / f"slide_{idx_int:03d}.png"
+            if path.exists():
+                return path.read_bytes()
+
+    return None
+
+
+def _select_slide_by_index(slides: list[dict], slide_index: int) -> dict | None:
+    if slide_index < 0:
+        return None
+    if slide_index >= len(slides):
+        return None
+    return slides[slide_index]
+
+
+def _gemini_translate_slide_image(client, image_bytes: bytes, language: str) -> tuple[bytes, str]:
+    prompt_language = (language or "").strip()
+    if not prompt_language:
+        raise ValueError("language must be a non-empty string")
+
+    with Image.open(io.BytesIO(image_bytes)) as source_image:
+        source_image.load()
+        inference_input = source_image.copy()
+
+    response = client.models.generate_content(
+        model="gemini-3-pro-image-preview",
+        contents=[f"Make a {prompt_language} version of this slide", inference_input],
+        config={
+            "response_modalities": ["IMAGE"],
+        },
+    )
+
+    parts = list(getattr(response, "parts", []) or [])
+    if not parts:
+        candidates = getattr(response, "candidates", None)
+        if candidates:
+            for candidate in candidates:
+                content = getattr(candidate, "content", None)
+                if content and getattr(content, "parts", None):
+                    parts.extend(content.parts)
+
+    for part in parts:
+        inline = getattr(part, "inline_data", None)
+        if inline:
+            data = getattr(inline, "data", None)
+            if data is None:
+                continue
+            if isinstance(data, str):
+                try:
+                    payload = base64.b64decode(data)
+                except Exception:
+                    continue
+            else:
+                payload = data
+            mime = getattr(inline, "mime_type", None) or "image/png"
+            return payload, mime
+
+    raise RuntimeError("Gemini did not return image data for the translated slide")
+
+
+def _translation_cache_paths(reference: str, language: str, slide_index: int) -> tuple[Path, Path]:
+    slug = _language_slug(language)
+    safe_index = max(0, int(slide_index))
+    base_dir = SLIDE_CACHE / reference / "translations" / slug
+    metadata_path = base_dir / f"slide_{safe_index:03d}.json"
+    return base_dir, metadata_path
+
+
+def _extension_for_mime(mime_type: str) -> str:
+    mapping = {
+        "image/png": "png",
+        "image/jpeg": "jpg",
+        "image/jpg": "jpg",
+        "image/webp": "webp",
+    }
+    mime = (mime_type or "").lower()
+    return mapping.get(mime, "bin")
+
+
 # ---------------------------------------------------------------------------------------------------------------------
 # Slide extraction pipeline
 # ---------------------------------------------------------------------------------------------------------------------
@@ -916,6 +1018,121 @@ def register_media_tools(app: FastMCP) -> None:
 
         return {"status": "not_found", "reference": reference}
 
+    @app.tool()
+    async def translate_slide(
+        ctx: Context,
+        reference: str,
+        slide_index: int,
+        language: str,
+    ) -> ImageContent:
+        """
+        Translate a previously extracted slide into another language using Gemini image-to-image.
+
+        Designed to be called after `start_media_retrieval` and `get_extracted_slides`.
+
+        Parameters:
+            - reference: Token returned by `start_media_retrieval` identifying the source media.
+            - slide_index: Zero-based slide number from `get_extracted_slides.slides[].index`.
+            - language: Target language name. Example: German.
+
+        Returns:
+            - image
+
+        Errors:
+            - All validation or runtime failures return `{ "is_error": true, "detail": "...", "reference": ... }`.
+        """
+
+        metadata = _load_json(_metadata_path(reference))
+        if not metadata or not Path(metadata.get("download_path", "")).exists():
+            return _error("media not downloaded", reference)
+
+        language_clean = (language or "").strip()
+        if not language_clean:
+            return _error("language must be provided", reference)
+
+        try:
+            slide_idx = int(slide_index)
+        except (TypeError, ValueError):
+            return _error("slide_index must be an integer", reference)
+        if slide_idx < 0:
+            return _error("slide_index must be >= 0", reference)
+
+        slides_payload = _load_json(_slides_json_path(reference))
+        if not slides_payload or not (slides_payload.get("slides") or []):
+            slides_payload = await asyncio.to_thread(_extract_slides_flow, metadata)
+
+        slides = slides_payload.get("slides") or []
+        if not slides:
+            return _error("no slides available for translation", reference)
+
+        slide = _select_slide_by_index(slides, slide_idx)
+        if not slide:
+            return _error("no slide matches the requested slide_index", reference)
+
+        slide_bytes = _slide_image_bytes(reference, slide)
+        if not slide_bytes:
+            return _error("slide image data missing", reference)
+
+        base_dir, metadata_path = _translation_cache_paths(reference, language_clean, slide_idx)
+        cached = False
+        translated_bytes: bytes | None = None
+        mime_type: str | None = None
+        dest_path: Path | None = None
+
+        if metadata_path.exists():
+            try:
+                record = json.loads(metadata_path.read_text())
+                filename = record.get("filename")
+                if filename:
+                    candidate = base_dir / filename
+                    if candidate.exists():
+                        translated_bytes = candidate.read_bytes()
+                        mime_type = record.get("mime_type") or "application/octet-stream"
+                        dest_path = candidate
+                        cached = True
+            except Exception:
+                pass
+
+        if translated_bytes is None or mime_type is None:
+            client = _build_gemini_client()
+            translated_bytes, mime_type = await asyncio.to_thread(
+                _gemini_translate_slide_image, client, slide_bytes, language_clean
+            )
+            mime_type = mime_type or "application/octet-stream"
+            extension = _extension_for_mime(mime_type)
+            image_filename = f"slide_{slide_idx:03d}.{extension}"
+            dest_path = base_dir / image_filename
+            dest_path.parent.mkdir(parents=True, exist_ok=True)
+            dest_path.write_bytes(translated_bytes)
+            metadata = {"mime_type": mime_type, "filename": image_filename}
+            metadata_path.write_text(json.dumps(metadata, indent=2))
+
+        mime_type = mime_type or "application/octet-stream"
+        if dest_path is None:
+            extension = _extension_for_mime(mime_type)
+            dest_path = base_dir / f"slide_{slide_idx:03d}.{extension}"
+            dest_path.parent.mkdir(parents=True, exist_ok=True)
+            dest_path.write_bytes(translated_bytes)
+        base64_data = base64.b64encode(translated_bytes).decode("ascii")
+        data_uri = f"data:{mime_type};base64,{base64_data}"
+        _write_debug(
+            reference,
+            f"translation_{_language_slug(language_clean)}_slide_{slide_idx:03d}.json",
+            {
+                "language": language_clean,
+                "slide_index": slide_idx,
+                "mime_type": mime_type,
+                "cached": cached,
+                "output_path": str(dest_path),
+            },
+        )
+
+        timestamp_mid = slide.get("mid")
+        timestamp_from = slide.get("from")
+        timestamp_to = slide.get("to")
+
+        return ImageContent(type="image", data=base64_data, mimeType=mime_type)
+
     @app.tool()
     async def start_media_analysis(
         ctx: Context,