cmhub/apps/ai/importers.py

from __future__ import annotations

from typing import Any, Mapping

from apps.ai.providers.utils import (
    API_CHAT,
    API_GEMINI,
    API_IMAGES,
    API_IMAGES_EDITS,
    detect_api_type,
)

from .models import AiModel, ModelAlias


def import_ai_models_config(
    config: Mapping[str, Any],
    *,
    create_default_aliases: bool = False,
) -> dict[str, int]:
    """Import cmbot-style ai_models.json data into AiModel records."""
    raw_models = config.get("models")
    if not isinstance(raw_models, list):
        raise ValueError("ai_models config must contain a models list")

    created = 0
    updated = 0
    imported_models: list[AiModel] = []
    for raw_model in raw_models:
        if not isinstance(raw_model, Mapping):
            raise ValueError("each model config must be an object")
        ai_model, was_created = _upsert_ai_model(raw_model)
        imported_models.append(ai_model)
        if was_created:
            created += 1
        else:
            updated += 1

    aliases = 0
    if create_default_aliases:
        aliases = _ensure_default_aliases(imported_models)

    return {"created": created, "updated": updated, "aliases": aliases}


def _upsert_ai_model(raw_model: Mapping[str, Any]) -> tuple[AiModel, bool]:
    name = str(raw_model.get("name") or raw_model.get("model") or "").strip()
    if not name:
        raise ValueError("model config is missing name/model")

    defaults = {
        "url": str(raw_model.get("url") or ""),
        "model": str(raw_model.get("model") or ""),
        "api_type": str(raw_model.get("api_type") or AiModel.ApiType.AUTO),
        "timeout_seconds": _to_non_negative_int(raw_model.get("timeout_seconds"), 0),
        "connect_timeout_seconds": _to_positive_int(
            raw_model.get("connect_timeout_seconds"), 30
        ),
        "extra_body": dict(raw_model.get("extra_body") or {}),
        "capabilities": _infer_capabilities(raw_model),
        "is_active": True,
    }
    ai_model, created = AiModel.objects.update_or_create(name=name, defaults=defaults)

    api_key = str(raw_model.get("api_key") or "").strip()
    if api_key:
        ai_model.set_api_key(api_key)
        ai_model.save(update_fields=("api_key_encrypted", "updated_at"))
    return ai_model, created


def _infer_capabilities(raw_model: Mapping[str, Any]) -> list[str]:
    explicit = raw_model.get("capabilities")
    if explicit:
        if isinstance(explicit, dict):
            return [
                str(key)
                for key, value in explicit.items()
                if key in {"text", "image", "vision"} and bool(value)
            ]
        if isinstance(explicit, str):
            return [item.strip() for item in explicit.split(",") if item.strip()]
        return [str(item) for item in explicit]

    api_type = str(raw_model.get("api_type") or AiModel.ApiType.AUTO)
    url = str(raw_model.get("url") or "")
    resolved_api_type = detect_api_type(url, api_type)
    name_model_text = f"{raw_model.get('name', '')} {raw_model.get('model', '')}".lower()

    if resolved_api_type == API_IMAGES_EDITS:
        return ["image", "vision"]
    if resolved_api_type == API_IMAGES:
        return ["image"]
    if resolved_api_type == API_GEMINI:
        return ["text", "image", "vision"]
    if resolved_api_type == API_CHAT and (
        "image" in name_model_text or "banana" in name_model_text
    ):
        return ["image", "vision"]
    return ["text", "vision"]


def _ensure_default_aliases(models: list[AiModel]) -> int:
    aliases = 0
    text_model = _first_model_with_capability(models, "text")
    image_model = _first_model_with_capability(models, "image")

    if text_model:
        ModelAlias.objects.filter(
            operation_type=ModelAlias.OperationType.TITLE,
            is_default=True,
        ).update(is_default=False)
        ModelAlias.objects.update_or_create(
            operation_type=ModelAlias.OperationType.TITLE,
            alias="title-standard",
            defaults={"ai_model": text_model, "is_default": True, "is_active": True},
        )
        aliases += 1

    if image_model:
        ModelAlias.objects.filter(
            operation_type=ModelAlias.OperationType.IMAGE,
            is_default=True,
        ).update(is_default=False)
        ModelAlias.objects.update_or_create(
            operation_type=ModelAlias.OperationType.IMAGE,
            alias="image-standard",
            defaults={"ai_model": image_model, "is_default": True, "is_active": True},
        )
        aliases += 1

    return aliases


def _first_model_with_capability(models: list[AiModel], capability: str) -> AiModel | None:
    for model in models:
        if capability in model.capabilities_set():
            return model
    return None


def _to_non_negative_int(value: Any, default: int) -> int:
    parsed = _to_int(value, default)
    return parsed if parsed >= 0 else default


def _to_positive_int(value: Any, default: int) -> int:
    parsed = _to_int(value, default)
    return parsed if parsed > 0 else default


def _to_int(value: Any, default: int) -> int:
    if value is None or value == "":
        return default
    try:
        return int(value)
    except (TypeError, ValueError):
        return default