Agentic-OS/backend/app/services/model_discovery.py

"""Live model discovery from Ollama, Gemini, DeepSeek, and OpenRouter APIs."""
from __future__ import annotations

import asyncio
import logging
from typing import Any

import httpx

from app.config import settings

logger = logging.getLogger(__name__)

# DeepSeek catalog fallbacks when /models is unavailable or incomplete.
KNOWN_DEEPSEEK_MODELS = [
    {"id": "deepseek-chat", "label": "DeepSeek Chat", "tier_hint": "economy"},
    {"id": "deepseek-reasoner", "label": "DeepSeek Reasoner", "tier_hint": "premium"},
]


async def list_ollama_models() -> dict[str, Any]:
    base = settings.ollama_base_url.rstrip("/")
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.get(f"{base}/api/tags")
            resp.raise_for_status()
            payload = resp.json()
    except Exception as exc:  # noqa: BLE001
        logger.debug("Ollama model list failed: %s", exc)
        return {"backend": "ollama", "available": False, "error": str(exc), "models": []}

    models = []
    for item in payload.get("models") or []:
        name = item.get("name") or item.get("model") or ""
        details = item.get("details") or {}
        caps = item.get("capabilities") or []
        models.append(
            {
                "id": name,
                "label": name,
                "size_bytes": item.get("size"),
                "parameter_size": details.get("parameter_size"),
                "quantization": details.get("quantization_level"),
                "context_length": details.get("context_length"),
                "capabilities": caps,
                "supports_tools": "tools" in caps,
            }
        )
    return {"backend": "ollama", "available": True, "models": models}


async def list_gemini_models() -> dict[str, Any]:
    if not settings.gemini_api_key:
        return {
            "backend": "gemini",
            "available": False,
            "error": "GEMINI_API_KEY not set",
            "models": [],
        }
    url = "https://generativelanguage.googleapis.com/v1beta/models"
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            resp = await client.get(url, params={"key": settings.gemini_api_key})
            resp.raise_for_status()
            payload = resp.json()
    except Exception as exc:  # noqa: BLE001
        logger.debug("Gemini model list failed: %s", exc)
        return {"backend": "gemini", "available": False, "error": str(exc), "models": []}

    models = []
    for item in payload.get("models") or []:
        methods = item.get("supportedGenerationMethods") or []
        if "generateContent" not in methods:
            continue
        raw_name = item.get("name") or ""
        model_id = raw_name.removeprefix("models/")
        display = item.get("displayName") or model_id
        models.append(
            {
                "id": model_id,
                "label": display,
                "input_token_limit": item.get("inputTokenLimit"),
                "output_token_limit": item.get("outputTokenLimit"),
                "description": item.get("description"),
            }
        )
    models.sort(key=lambda m: m["id"])
    return {"backend": "gemini", "available": True, "models": models}


async def list_deepseek_models() -> dict[str, Any]:
    if not settings.deepseek_api_key:
        return {
            "backend": "deepseek",
            "available": False,
            "error": "DEEPSEEK_API_KEY not set",
            "models": [],
        }
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            resp = await client.get(
                "https://api.deepseek.com/models",
                headers={"Authorization": f"Bearer {settings.deepseek_api_key}"},
            )
            resp.raise_for_status()
            payload = resp.json()
        models = [
            {
                "id": item.get("id") or "",
                "label": item.get("id") or "",
                "owned_by": item.get("owned_by"),
            }
            for item in payload.get("data") or []
            if item.get("id")
        ]
        if not models:
            models = [dict(m) for m in KNOWN_DEEPSEEK_MODELS]
        return {"backend": "deepseek", "available": True, "models": models}
    except Exception as exc:  # noqa: BLE001
        logger.debug("DeepSeek model list failed: %s", exc)
        return {
            "backend": "deepseek",
            "available": True,
            "error": str(exc),
            "models": [dict(m) for m in KNOWN_DEEPSEEK_MODELS],
            "note": "Using known DeepSeek model IDs (API list unavailable)",
        }


async def list_openrouter_models() -> dict[str, Any]:
    if not settings.openrouter_api_key:
        return {
            "backend": "openrouter",
            "available": False,
            "error": "OPENROUTER_API_KEY not set",
            "models": [],
        }
    try:
        async with httpx.AsyncClient(timeout=20.0) as client:
            resp = await client.get(
                "https://openrouter.ai/api/v1/models",
                headers={"Authorization": f"Bearer {settings.openrouter_api_key}"},
            )
            resp.raise_for_status()
            payload = resp.json()
    except Exception as exc:  # noqa: BLE001
        logger.debug("OpenRouter model list failed: %s", exc)
        return {"backend": "openrouter", "available": False, "error": str(exc), "models": []}

    models = []
    for item in payload.get("data") or []:
        pricing = item.get("pricing") or {}
        models.append(
            {
                "id": item.get("id") or "",
                "label": item.get("name") or item.get("id") or "",
                "context_length": item.get("context_length"),
                "input_per_token": pricing.get("prompt"),
                "output_per_token": pricing.get("completion"),
                "description": item.get("description"),
            }
        )
    models.sort(key=lambda m: m["id"])
    return {"backend": "openrouter", "available": True, "models": models}


async def list_all_provider_models() -> dict[str, Any]:
    ollama, gemini, deepseek, openrouter = await asyncio.gather(
        list_ollama_models(),
        list_gemini_models(),
        list_deepseek_models(),
        list_openrouter_models(),
    )
    return {
        "providers": {
            "ollama": ollama,
            "gemini": gemini,
            "deepseek": deepseek,
            "openrouter": openrouter,
        },
        "routing_order": settings.llm_cloud_provider_order,
    }