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-hd", 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