Moeblack-ComfyUI-SimpleChat/api/routes.py

"""
API routes for fetching model lists from various providers.
"""
import json
import aiohttp
from aiohttp import web

try:
    from server import PromptServer
    HAS_SERVER = True
except ImportError:
    HAS_SERVER = False


# Predefined model lists for each provider
PREDEFINED_MODELS = {
    "openai": [
        "gpt-4o",
        "gpt-4o-mini",
        "gpt-4-turbo",
        "gpt-4",
        "gpt-3.5-turbo",
        "o1-preview",
        "o1-mini",
    ],
    "claude": [
        "claude-sonnet-4-20250514",
        "claude-3-5-sonnet-20241022",
        "claude-3-5-haiku-20241022",
        "claude-3-opus-20240229",
        "claude-3-sonnet-20240229",
        "claude-3-haiku-20240307",
    ],
    "gemini": [
        "gemini-2.5-flash-preview-05-20",
        "gemini-2.0-flash",
        "gemini-2.0-flash-exp",
        "gemini-1.5-pro",
        "gemini-1.5-flash",
        "gemini-2.5-flash-image",
        "gemini-3-pro-image-preview",
    ],
}


async def fetch_openai_models(api_key: str, base_url: str) -> list[str]:
    """Fetch model list from OpenAI-compatible API."""
    url = f"{base_url.rstrip('/')}/models"
    headers = {"Authorization": f"Bearer {api_key}"}

    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, headers=headers, timeout=10) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    models = [m["id"] for m in data.get("data", [])]
                    # Return all models, sorted alphabetically
                    return sorted(models) if models else PREDEFINED_MODELS["openai"]
    except Exception as e:
        print(f"[SimpleChat] Failed to fetch OpenAI models: {e}")

    return PREDEFINED_MODELS["openai"]


async def fetch_claude_models(api_key: str, base_url: str) -> list[str]:
    """
    Fetch models for Claude provider.
    Tries to hit the /models endpoint (supported by some proxies and newer Anthropic API).
    Falls back to predefined list if fails.
    """
    # Anthropic API format for models list is /v1/models
    # Note: Official Anthropic API recently added models endpoint support
    url = f"{base_url.rstrip('/')}/models"
    headers = {
        "x-api-key": api_key,
        "anthropic-version": "2023-06-01"
    }

    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, headers=headers, timeout=10) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    # Anthropic format: {"data": [{"id": "...", ...}]} similar to OpenAI
                    # Some proxies might return OpenAI format
                    models = []

                    # Handle both standard formats just in case
                    items = data.get("data", []) if "data" in data else data.get("models", [])

                    if isinstance(items, list):
                        for m in items:
                            if isinstance(m, dict):
                                models.append(m.get("id", m.get("name", "")))
                            elif isinstance(m, str):
                                models.append(m)

                    # Filter out empty strings
                    models = [m for m in models if m]

                    if models:
                        return sorted(models)
    except Exception as e:
        # Silently fail for Claude since it's common for this endpoint to not exist/fail
        pass

    return PREDEFINED_MODELS["claude"]


async def fetch_gemini_models(api_key: str, base_url: str) -> list[str]:
    """Fetch model list from Gemini API."""
    url = f"{base_url.rstrip('/')}/models?key={api_key}"

    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, timeout=10) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    models = []
                    for m in data.get("models", []):
                        name = m.get("name", "")
                        # Extract model name from full path (models/gemini-xxx)
                        if name.startswith("models/"):
                            name = name[7:]
                        if name:
                            models.append(name)
                    # Return all models
                    return sorted(models) if models else PREDEFINED_MODELS["gemini"]
    except Exception as e:
        print(f"[SimpleChat] Failed to fetch Gemini models: {e}")

    return PREDEFINED_MODELS["gemini"]


def get_all_predefined_models_list() -> list[str]:
    """Get a flat list of all predefined models."""
    all_models = []
    for provider_models in PREDEFINED_MODELS.values():
        all_models.extend(provider_models)
    return all_models


def setup_routes():
    """Setup API routes for SimpleChat."""
    if not HAS_SERVER:
        return

    @PromptServer.instance.routes.get("/simplechat/models/{provider}")
    async def get_models(request):
        """
        Get available models for a provider.

        Query params:
            - api_key: API key for the provider
            - base_url: (optional) Custom base URL
        """
        provider = request.match_info.get("provider", "").lower()
        api_key = request.query.get("api_key", "")
        base_url = request.query.get("base_url", "")

        # Return predefined models if no API key provided
        if not api_key:
            models = PREDEFINED_MODELS.get(provider, [])
            return web.json_response(models)

        # Default base URLs
        default_urls = {
            "openai": "https://api.openai.com/v1",
            "claude": "https://api.anthropic.com/v1",
            "gemini": "https://generativelanguage.googleapis.com/v1beta",
        }

        if not base_url:
            base_url = default_urls.get(provider, "")

        # Fetch models based on provider
        if provider == "openai":
            models = await fetch_openai_models(api_key, base_url)
        elif provider == "claude":
            models = await fetch_claude_models(api_key, base_url)
        elif provider == "gemini":
            models = await fetch_gemini_models(api_key, base_url)
        else:
            models = []

        return web.json_response(models)

    @PromptServer.instance.routes.get("/simplechat/models")
    async def get_all_models_list(request):
        """
        Get all predefined models as a flat list (for COMBO remote input).
        Returns a simple array of model names.
        """
        return web.json_response(get_all_predefined_models_list())

    @PromptServer.instance.routes.get("/simplechat/models_dict")
    async def get_all_predefined_models(request):
        """Get all predefined models organized by provider."""
        return web.json_response(PREDEFINED_MODELS)

    print("[SimpleChat] API routes registered")