live-radio/src/ozan_radio/lyria_capabilities.py

"""Lyria 3 capabilities — synced from Google Gemini API docs + live model list."""

from __future__ import annotations

from dataclasses import dataclass

from google import genai
from google.genai import errors as genai_errors

DOCS_URL = "https://ai.google.dev/gemini-api/docs/music-generation"

VOCAL_MODES = [
    {
        "id": "instrumental",
        "label": "Instrumental only",
        "hint": "Lyria tip: include 'Instrumental only, no vocals' in every prompt.",
    },
    {
        "id": "mix",
        "label": "Mix (textures + chants)",
        "hint": "Wordless vocals, Sahel chants, whispers — no full lead lyrics unless chat asks.",
    },
    {
        "id": "vocals",
        "label": "Full vocals",
        "hint": "Lead vocals and lyrics; language follows prompt or Lyria language setting.",
    },
]

LANGUAGES = [
    {"id": "auto", "label": "Auto (match prompt / taste)"},
    {"id": "en", "label": "English"},
    {"id": "hi", "label": "Hindi"},
    {"id": "tr", "label": "Turkish"},
    {"id": "fr", "label": "French"},
    {"id": "ar", "label": "Arabic"},
    {"id": "es", "label": "Spanish"},
    {"id": "pt", "label": "Portuguese"},
    {"id": "de", "label": "German"},
]

SINGER_PROFILES = [
    {"id": "", "label": "Default (model picks)"},
    {"id": "male_baritone", "label": "Male baritone — warm, desert blues"},
    {"id": "male_tenor", "label": "Male tenor — bright, energetic"},
    {"id": "female_alto", "label": "Female alto — smoky, soulful"},
    {"id": "female_soprano", "label": "Female soprano — airy, soaring"},
    {"id": "weathered_rocker", "label": "Weathered rocker — raspy 90s grit"},
]

OUTPUT_FORMATS = [
    {"id": "mp3", "label": "MP3 (default)"},
    {"id": "wav", "label": "WAV (Pro only, larger files)"},
]

STATIC_MODELS = [
    {
        "id": "lyria-3-pro-preview",
        "label": "Lyria 3 Pro",
        "duration": "1–2 minutes",
        "cost_key": "lyria_pro_usd",
        "supports_wav": True,
        "supports_timestamps": True,
        "supports_image_input": True,
    },
    {
        "id": "lyria-3-clip-preview",
        "label": "Lyria 3 Clip",
        "duration": "30 seconds",
        "cost_key": "lyria_clip_usd",
        "supports_wav": False,
        "supports_timestamps": False,
        "supports_image_input": True,
    },
]


@dataclass
class LyriaApiStatus:
    ok: bool
    message: str
    models_available: list[str]
    key_configured: bool


def static_capabilities() -> dict:
    return {
        "docs_url": DOCS_URL,
        "models": STATIC_MODELS,
        "vocal_modes": VOCAL_MODES,
        "languages": LANGUAGES,
        "singer_profiles": SINGER_PROFILES,
        "output_formats": OUTPUT_FORMATS,
        "features": [
            "Text-to-music via generateContent",
            "Vocals + timed lyrics (Pro)",
            "Instrumental-only prompts",
            "Custom lyrics with [Verse]/[Chorus] tags",
            "Timestamp structure [0:00 - 0:30]",
            "Prompt language steers lyric language",
            "Image-to-music (multimodal)",
            "44.1 kHz stereo MP3 or WAV (Pro)",
        ],
        "prompt_tips": [
            "Iterate with Clip (30s) before burning a Pro generation.",
            "Be specific: BPM, key, instruments, mood, structure.",
            "For instrumentals, always say 'Instrumental only, no vocals'.",
            "For vocals, describe singer gender, timbre, and range.",
            "Avoid 'vocal sample' / 'griot sample' — use 'Sahel blues chants'.",
        ],
    }


def probe_lyria_api(api_key: str | None) -> dict:
    if not api_key:
        return LyriaApiStatus(
            ok=False,
            message="GEMINI_API_KEY not set",
            models_available=[],
            key_configured=False,
        ).__dict__

    try:
        client = genai.Client(api_key=api_key)
        found: list[str] = []
        for model in client.models.list():
            name = (getattr(model, "name", "") or "").replace("models/", "")
            if "lyria" in name.lower():
                found.append(name)
        if not found:
            return LyriaApiStatus(
                ok=False,
                message="API key works but no Lyria models visible on this project",
                models_available=[],
                key_configured=True,
            ).__dict__
        return LyriaApiStatus(
            ok=True,
            message=f"Lyria ready — {len(found)} model(s) available",
            models_available=sorted(found),
            key_configured=True,
        ).__dict__
    except genai_errors.ClientError as exc:
        return LyriaApiStatus(
            ok=False,
            message=f"Gemini API error: {exc}",
            models_available=[],
            key_configured=True,
        ).__dict__
    except Exception as exc:
        return LyriaApiStatus(
            ok=False,
            message=f"Probe failed: {exc}",
            models_available=[],
            key_configured=True,
        ).__dict__