ai_v/utils.py

from urllib.parse import quote, urlparse, urlunparse
from config import Config

def get_proxied_url(url):
    """
    如果启用了代理，则返回经过代理包装的 URL。
    否则返回原 URL。
    """
    if Config.USE_PROXY and url:
        return f"{Config.PROXY_URL}{quote(url)}"
    return url

def get_backup_api_url(url):
    """
    将已有接口地址切换到备用 AI Base URL，并保留原始路径。
    """
    if not url or not Config.AI_BACKUP_BASE_URL:
        return url

    source = urlparse(url)
    backup = urlparse(Config.AI_BACKUP_BASE_URL)
    if not backup.scheme or not backup.netloc:
        return url

    return urlunparse((
        backup.scheme,
        backup.netloc,
        source.path,
        source.params,
        source.query,
        source.fragment,
    ))

def get_backup_api_key(api_key):
    """
    返回内置主 Key 对应的备用 Key；自定义 Key 默认沿用原值。
    """
    fallback_map = {
        Config.TRIAL_KEY: Config.TRIAL_BACKUP_KEY,
        Config.PREMIUM_KEY: Config.PREMIUM_BACKUP_KEY,
        Config.VIDEO_KEY: Config.VIDEO_BACKUP_KEY,
        Config.GEMINI_FLASH_IMAGE_PREMIUM_KEY: Config.GEMINI_FLASH_IMAGE_PREMIUM_BACKUP_KEY,
    }
    return fallback_map.get(api_key, api_key)

def get_api_candidates(url, api_key):
    """
    生成主线路与备用线路候选。
    """
    candidates = [{"url": url, "api_key": api_key, "label": "primary"}]
    backup_url = get_backup_api_url(url)
    backup_key = get_backup_api_key(api_key)

    if backup_url and backup_url != url:
        candidates.append({"url": backup_url, "api_key": backup_key, "label": "backup"})

    return candidates

def should_switch_to_backup(status_code):
    """
    遇到鉴权、限流或服务端异常时，尝试切换备用线路。
    """
    return status_code in (401, 403, 408, 429) or status_code >= 500