src/models.py

from __future__ import annotations

from typing import Any

from src.config import settings

# Intentionally small registry. This is not an arbitrary repo downloader.
MODEL_REGISTRY: dict[str, dict[str, Any]] = {
    settings.model_name: {
        "id": settings.model_name,
        "repo_id": settings.gguf_repo_id,
        "filename": settings.gguf_filename,
        "owned_by": "unsloth",
        "backend": "llama-diffusion-cli",
        "supports_openai_stream": "final_delta_only",
        "supports_diffusion_visual": True,
    },
    # Convenience aliases.
    settings.gguf_repo_id: {
        "id": settings.model_name,
        "repo_id": settings.gguf_repo_id,
        "filename": settings.gguf_filename,
        "owned_by": "unsloth",
        "backend": "llama-diffusion-cli",
        "supports_openai_stream": "final_delta_only",
        "supports_diffusion_visual": True,
    },
}


def default_model_id() -> str:
    return settings.model_name


def get_model_config(model_id: str | None) -> dict[str, Any]:
    requested = model_id or default_model_id()
    if requested not in MODEL_REGISTRY:
        available = sorted({v["id"] for v in MODEL_REGISTRY.values()})
        from src.errors import ApiError
        raise ApiError(
            code="unsupported_model",
            message=f"Unsupported model: {requested}. Available models: {available}",
            status_code=400,
        )
    cfg = dict(MODEL_REGISTRY[requested])
    cfg["requested_id"] = requested
    return cfg


def public_model_list() -> list[dict[str, Any]]:
    seen: set[str] = set()
    models: list[dict[str, Any]] = []
    for cfg in MODEL_REGISTRY.values():
        mid = cfg["id"]
        if mid in seen:
            continue
        seen.add(mid)
        models.append(
            {
                "id": mid,
                "object": "model",
                "owned_by": cfg.get("owned_by", "unsloth"),
                "backend": cfg.get("backend"),
            }
        )
    return models