tts-voice-creator-clone-and-design-2/core/provider_health.py

"""Shared, bounded provider liveness probes."""
from __future__ import annotations

from typing import Any, Iterable
from urllib.parse import urlsplit

import requests


def probe_provider(
    base_url: str,
    checks: Iterable[tuple[str, str | None]],
    *,
    headers: dict[str, str] | None = None,
    timeout: float = 3.0,
) -> dict[str, Any]:
    """Try provider endpoints in order with normalized, body-free results."""
    base = str(base_url or '').rstrip('/')
    limit = max(0.5, float(timeout))
    last_error = 'No response'
    for path, json_key in checks:
        endpoint = path or '/'
        try:
            response = requests.get(base + path, timeout=limit, headers=headers or {}, allow_redirects=True)
            if response.status_code >= 500:
                last_error = f'HTTP {response.status_code} on {endpoint}'
                continue
            if response.status_code >= 400:
                last_error = f'HTTP {response.status_code} on {endpoint}'
                continue
            payload = None
            if json_key:
                try:
                    payload = response.json()
                except Exception:
                    last_error = f'{endpoint} returned non-JSON (HTTP {response.status_code})'
                    continue
                if not isinstance(payload, dict) or json_key not in payload:
                    last_error = f"Unexpected response from {endpoint} (missing '{json_key}')"
                    continue
            return {'ok': True, 'status': response.status_code, 'endpoint': endpoint, 'payload': payload}
        except requests.exceptions.ConnectionError:
            last_error = 'Connection refused'
        except requests.exceptions.Timeout:
            last_error = f'Timeout ({limit:g} s)'
        except Exception as exc:
            last_error = str(exc)
    return {'ok': False, 'error': last_error}


def provider_key_for_url(settings: dict, url: str, fallback: str = "") -> str:
    """Resolve saved card credentials for the endpoint actually being called."""
    normalized = str(url).rstrip("/").removesuffix("/v1")
    cards = dict(settings.get("engine_local_urls") or {})
    for card in settings.get("custom_engine_cards") or []:
        cards[str(card.get("id") or card.get("name") or "")] = card.get("url", "")
    keys = settings.get("engine_api_keys") or {}
    for name, card_url in cards.items():
        if str(card_url).rstrip("/").removesuffix("/v1") == normalized:
            return str(keys.get(name.removeprefix("dc-")) or keys.get(name) or "")
    return fallback


def _litellm_model_aliases(base: str, key: str) -> list[str]:
    """Names of a LiteLLM gateway's router aliases (Model Group Alias), which its /models list omits.

    Best effort: the settings endpoint needs an admin-capable key, and other gateways do not have it.
    """
    root = base[:-3] if base.endswith("/v1") else base
    try:
        r = requests.get(f"{root}/router/settings", timeout=5, headers={"Authorization": f"Bearer {key}"})
        if r.status_code != 200:
            return []
        data = r.json()
    except Exception:
        return []
    if not isinstance(data, dict):
        return []
    for holder in (data.get("current_values"), data.get("router_settings"), data):
        aliases = holder.get("model_group_alias") if isinstance(holder, dict) else None
        if isinstance(aliases, dict):
            return [str(name) for name in aliases if name]
    return []


_PROVIDER_REQUIRED_FIELDS = {
    'model', 'messages', 'input', 'text', 'voice', 'stream', 'response_format',
    'file', 'files', 'audio', 'audio_prompt', 'audio_prompt_transcript', 'references',
    'url', 'base_url', 'endpoint', 'headers', 'method',
}
_PROVIDER_SECRET_FIELDS = {
    'apikey', 'authorization', 'proxyauthorization', 'password', 'passwd', 'secret',
    'token', 'accesstoken', 'refreshtoken', 'privatekey', 'credentials', 'credential',
    'xapikey', 'accesskey', 'secretkey',
}


def validate_provider_params(params: Any) -> None:
    """Keep credentials and required request fields out of optional metadata."""
    if not isinstance(params, dict):
        raise ValueError('Additional parameters must be a JSON object')
    if any(key in _PROVIDER_REQUIRED_FIELDS for key in params):
        raise ValueError('Additional parameters cannot replace required request fields')

    def check(value):
        if isinstance(value, dict):
            for key, item in value.items():
                normalized = ''.join(character for character in str(key).lower() if character.isalnum())
                if normalized in _PROVIDER_SECRET_FIELDS or normalized == 'headers' or normalized.endswith(('apikey', 'authorization', 'password', 'secret', 'token', 'secretkey', 'privatekey')):
                    raise ValueError('Credentials must use the API key field')
                check(item)
        elif isinstance(value, list):
            for item in value:
                check(item)
    check(params)


def engine_disconnected(settings: dict, url: str, role: str = "", identity: str = "") -> bool:
    """Use explicit selection first; infer legacy endpoints only when unambiguous."""
    base = str(url).rstrip('/').removesuffix('/v1')
    urls = dict(settings.get('engine_local_urls') or {})
    urls.update({str(name).removeprefix('dc-'): value for name, value in list(urls.items()) if str(name).startswith('dc-')})
    urls = {name: value for name, value in urls.items() if not str(name).startswith('dc-')}
    from routes.docker import _LOCAL_CONTAINER_DEFS
    roles = {card['name']: card['role'] for card in _LOCAL_CONTAINER_DEFS}
    parsed = urlsplit(base)
    if parsed.hostname in {'localhost', '127.0.0.1', '::1', 'host.docker.internal'}:
        for card in _LOCAL_CONTAINER_DEFS:
            if card['port'] == parsed.port:
                urls.setdefault(card['name'], base)
    for card in settings.get('custom_engine_cards') or []:
        name = str(card.get('id') or card.get('name') or '')
        roles[name] = str(card.get('role') or '')
        urls.setdefault(name, card.get('url'))
    if not identity:
        for setting, selected in (settings.get('engine_active') or {}).items():
            if str(settings.get(setting) or '').rstrip('/').removesuffix('/v1') == base:
                if setting == {'llm': 'llm_url', 'ocr': 'vision_ocr_url'}.get(role):
                    selected_url = urls.get(selected)
                    if not selected_url or str(selected_url).rstrip('/').removesuffix('/v1') == base:
                        identity = selected
                        break
    if not identity:
        matches = [name for name, value in urls.items() if (not role or not roles.get(name) or role in roles[name].split('+')) and str(value).rstrip('/').removesuffix('/v1') == base]
        if len(matches) == 1:
            identity = matches[0]
    return (settings.get('engine_disabled') or {}).get(identity) is True


def provider_request_params(url: str, payload: dict, settings: dict | None = None, role: str = "", identity: str = "") -> dict:
    """Merge custom endpoint options while preserving required request fields."""
    if settings is None:
        from core.config import _load_settings
        settings = _load_settings()
    base = str(url).rstrip('/')
    routes = {'/audio/transcriptions': 'stt', '/audio/speech': 'tts', '/chat/completions': 'llm', '/transcribe': 'stt', '/tts': 'tts'}
    for suffix, route_role in routes.items():
        if base.endswith(suffix):
            base = base[:-len(suffix)]
            role = role or route_role
            break
    base = base.removesuffix('/v1')
    if engine_disconnected(settings, base, role, identity):
        raise ValueError('Engine is disconnected; connect it before use')
    candidates = []
    urls = settings.get('engine_local_urls') or {}
    for card in settings.get('custom_engine_cards') or []:
        card_role = str(card.get('role') or '')
        if role and card_role and role not in card_role.split('+'):
            continue
        card_identity = str(card.get('id') or card.get('name') or '')
        if identity and identity != card_identity:
            continue
        legacy_name = str(card.get('name') or '')
        configured = urls.get('dc-' + card_identity) or urls.get(card_identity) or urls.get('dc-' + legacy_name) or urls.get(legacy_name) or card.get('url') or ''
        if str(configured).rstrip('/').removesuffix('/v1') == base:
            candidates.append(card)
    if len(candidates) != 1:
        return payload
    params = candidates[0].get('params') or {}
    try:
        validate_provider_params(params)
    except ValueError:
        return payload
    return {**payload, **params, **{key: value for key, value in payload.items() if key in _PROVIDER_REQUIRED_FIELDS}}