backend/modal_client.py

import base64
import os
import time
from pathlib import Path
from typing import Dict, Iterable, List, Optional
from urllib.parse import urlparse

import requests

from backend.synthesis_catalog import MODAL_BACKEND
from backend.types import VoiceConfig


class ModalSynthesisClient:
    def __init__(
        self,
        *,
        base_url: Optional[str],
        auth_token: Optional[str] = None,
        timeout_seconds: float = 300.0,
        poll_interval_seconds: float = 3.0,
        max_status_retries: int = 4,
    ) -> None:
        self.base_url = base_url.rstrip("/") if base_url else None
        self.auth_token = auth_token
        self.timeout_seconds = timeout_seconds
        self.poll_interval_seconds = poll_interval_seconds
        self.max_status_retries = max_status_retries

    @classmethod
    def from_env(cls) -> "ModalSynthesisClient":
        return cls(
            base_url=os.getenv("SCRIPTORIUM_MODAL_BASE_URL"),
            auth_token=os.getenv("SCRIPTORIUM_MODAL_AUTH_TOKEN"),
            timeout_seconds=float(os.getenv("SCRIPTORIUM_MODAL_TIMEOUT_SECONDS", "300")),
            poll_interval_seconds=float(os.getenv("SCRIPTORIUM_MODAL_POLL_INTERVAL_SECONDS", "3")),
            max_status_retries=int(os.getenv("SCRIPTORIUM_MODAL_STATUS_RETRIES", "4")),
        )

    def is_configured(self) -> bool:
        return bool(self.base_url)

    def configuration_warning(self) -> Optional[str]:
        if not self.base_url:
            return None

        parsed = urlparse(self.base_url)
        host = (parsed.netloc or "").lower()
        path = parsed.path or ""
        if host == "modal.com" and path.startswith("/apps/"):
            return (
                "SCRIPTORIUM_MODAL_BASE_URL points to a Modal dashboard URL, not a callable web endpoint. "
                "Use the deployed modal.run base URL instead, for example "
                "'https://<your-app>.modal.run'."
            )
        return None

    def generate_preview(
        self,
        *,
        text: str,
        output_path: Path,
        voice_config: VoiceConfig,
        diffusion_steps: int,
        speed: float,
    ) -> Dict[str, object]:
        payload = {
            "text": text,
            "voice_config": voice_config.to_dict(),
            "diffusion_steps": diffusion_steps,
            "speed": speed,
        }
        data = self._post_json("/preview", payload)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_bytes(base64.b64decode(data["audio_base64"]))
        return {
            "duration_seconds": int(data.get("duration_seconds", 0)),
            "sample_rate": int(data.get("sample_rate", 24000)),
            "backend": MODAL_BACKEND,
            "model": voice_config.model,
        }

    def submit_render(
        self,
        *,
        session_id: str,
        book: Dict[str, object],
        chapters: List[Dict[str, object]],
        voice_config: VoiceConfig,
        diffusion_steps: int,
        speed: float,
    ) -> str:
        payload = {
            "session_id": session_id,
            "book": book,
            "chapters": chapters,
            "voice_config": voice_config.to_dict(),
            "diffusion_steps": diffusion_steps,
            "speed": speed,
        }
        data = self._post_json("/renders", payload)
        job_id = data.get("job_id")
        if not job_id:
            raise ValueError("Modal render submission did not return a job id")
        return str(job_id)

    def render(
        self,
        *,
        session_id: str,
        job_id: str,
        render_dir: Path,
        voice_config: VoiceConfig,
    ) -> Iterable[Dict[str, object]]:
        cursor = 0
        consecutive_status_failures = 0
        while True:
            try:
                data = self._get_json(f"/renders/{job_id}", params={"cursor": cursor})
                consecutive_status_failures = 0
            except ValueError as exc:
                consecutive_status_failures += 1
                if consecutive_status_failures > self.max_status_retries:
                    raise
                yield {
                    "type": "log",
                    "session_id": session_id,
                    "backend": MODAL_BACKEND,
                    "model": voice_config.model,
                    "message": (
                        f"Transient Modal status error ({consecutive_status_failures}/{self.max_status_retries}): {exc}"
                    ),
                }
                time.sleep(self.poll_interval_seconds)
                continue
            events = list(data.get("events") or [])
            cursor += len(events)
            for event in events:
                yield self._materialize_event(
                    session_id=session_id,
                    render_dir=render_dir,
                    voice_config=voice_config,
                    event=event,
                )
            status = str(data.get("status", "pending"))
            if status in {"completed", "failed", "cancelled"}:
                break
            time.sleep(self.poll_interval_seconds)

    def cancel_render(self, job_id: str) -> Dict[str, object]:
        data = self._post_json(f"/renders/{job_id}/cancel", {})
        return {
            "type": data.get("type", "cancelled"),
            "job_id": job_id,
            "backend": MODAL_BACKEND,
        }

    def _materialize_event(
        self,
        *,
        session_id: str,
        render_dir: Path,
        voice_config: VoiceConfig,
        event: Dict[str, object],
    ) -> Dict[str, object]:
        materialized = dict(event)
        materialized.setdefault("session_id", session_id)
        materialized.setdefault("backend", MODAL_BACKEND)
        materialized.setdefault("model", voice_config.model)
        artifact_url = materialized.get("artifact_url")
        filename = materialized.get("filename")
        if artifact_url and filename:
            output_path = render_dir / str(filename)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            response = requests.get(
                self._absolute_url(str(artifact_url)),
                headers=self._headers(),
                timeout=self.timeout_seconds,
            )
            response.raise_for_status()
            output_path.write_bytes(response.content)
            materialized["output_path"] = str(output_path)
        return materialized

    def _post_json(self, path: str, payload: Dict[str, object]) -> Dict[str, object]:
        self._ensure_configured()
        try:
            response = requests.post(
                self._absolute_url(path),
                json=payload,
                headers=self._headers(),
                timeout=self.timeout_seconds,
            )
            response.raise_for_status()
            return response.json()
        except requests.Timeout as exc:
            raise ValueError(
                "Modal request timed out. The remote worker may still be cold-starting or loading models. "
                "Try again, or increase SCRIPTORIUM_MODAL_TIMEOUT_SECONDS."
            ) from exc
        except requests.RequestException as exc:
            detail = ""
            response = getattr(exc, "response", None)
            if response is not None and response.text:
                detail = f" | response={response.text[:400]}"
            raise ValueError(f"Modal request failed: {exc}{detail}") from exc

    def _get_json(self, path: str, *, params: Dict[str, object]) -> Dict[str, object]:
        self._ensure_configured()
        try:
            response = requests.get(
                self._absolute_url(path),
                params=params,
                headers=self._headers(),
                timeout=self.timeout_seconds,
            )
            response.raise_for_status()
            return response.json()
        except requests.Timeout as exc:
            raise ValueError(
                "Modal request timed out. The remote worker may still be cold-starting or loading models. "
                "Try again, or increase SCRIPTORIUM_MODAL_TIMEOUT_SECONDS."
            ) from exc
        except requests.RequestException as exc:
            detail = ""
            response = getattr(exc, "response", None)
            if response is not None and response.text:
                detail = f" | response={response.text[:400]}"
            raise ValueError(f"Modal request failed: {exc}{detail}") from exc

    def _absolute_url(self, path: str) -> str:
        if path.startswith("http://") or path.startswith("https://"):
            return path
        return f"{self.base_url}{path}"

    def _headers(self) -> Dict[str, str]:
        headers = {"Accept": "application/json"}
        if self.auth_token:
            headers["Authorization"] = f"Bearer {self.auth_token}"
        return headers

    def _ensure_configured(self) -> None:
        if not self.base_url:
            raise ValueError("Modal backend is not configured. Set SCRIPTORIUM_MODAL_BASE_URL.")