models/registry.py

"""
NeuraPrompt Model Registry
==========================
Auto-discovers and loads all MODEL dicts from this package.
Drop a new model .py file here and restart — it appears automatically.
"""

import importlib
import pathlib
import logging

_registry: dict[str, dict] = {}

DEFAULT_MODEL_ID = "neurones-pro-1.0"


def load_all():
    """Scan models/ and load every file that defines a MODULE-level MODEL dict."""
    models_dir = pathlib.Path(__file__).parent
    for f in sorted(models_dir.glob("*.py")):
        if f.name.startswith("_") or f.name == "registry.py":
            continue
        try:
            mod = importlib.import_module(f"models.{f.stem}")
            if hasattr(mod, "MODEL"):
                m = mod.MODEL
                _registry[m["id"]] = m
                logging.info(f"✅ Model registered: {m['display_name']} ({m['id']})")
        except Exception as e:
            logging.error(f"❌ Failed to load model {f.name}: {e}")
    logging.info(f"📦 {len(_registry)} model(s) loaded: {list(_registry.keys())}")


def get(model_id: str) -> dict:
    """Return model config by ID. Falls back to default if not found."""
    if model_id in _registry:
        return _registry[model_id]
    # Legacy AIModel enum value fallback mapping
    _legacy = {
        "neurones_self":       "neurones-pro-1.0",
        "neurones_self_3_0":   "neurones-pro-1.0",
        "neurones-r1-1.1":     "neurones-r1-1.1",
        "openai/gpt-oss-safeguard-20b": "neurones-flash-2.0",
        "openai/gpt-oss-120b":          "neurones-pro-1.0",
    }
    fallback_id = _legacy.get(model_id, DEFAULT_MODEL_ID)
    if fallback_id in _registry:
        return _registry[fallback_id]
    # Last resort: return whatever is first
    return next(iter(_registry.values())) if _registry else _get_emergency_fallback()


def default() -> dict:
    return _registry.get(DEFAULT_MODEL_ID) or next(iter(_registry.values()), _get_emergency_fallback())


def list_all() -> list[dict]:
    return list(_registry.values())


def is_loaded() -> bool:
    return len(_registry) > 0


def _get_emergency_fallback() -> dict:
    """Safety net if no models loaded at all."""
    return {
        "id": "neurones-pro-1.0",
        "display_name": "Neurones Pro 1.0",
        "version": "1.1",
        "tagline": "Default model",
        "speed": "fast",
        "speed_label": "🚀 Fast",
        "groq_model": "openai/gpt-oss-120b",
        "max_tokens": 4096,
        "temperature": 0.8,
        "can_stream": True,
        "can_reason": True,
        "can_vision": False,
        "can_generate_image": False,
        "can_search": True,
        "can_code": True,
        "can_translate": True,
        "can_summarise": True,
        "is_local": False,
        "context_window": 8192,
        "rate_limit_rpm": 50,
        "system_prompt": "You are NeuraPrompt AI, created by Andile Mtolo (Toxic Dee Modder). Be helpful and accurate.",
        "badge_color": "#00c853",
        "icon": "🚀",
        "recommended_for": ["general tasks"],
        "not_recommended_for": ["deep reasoning, image, file"],
    }