backend/services/models.py

from typing import Any

from fastapi import HTTPException

from models.llama_3_8b import Llama3_8B
from models.mistral_7b import Mistral_7b
from models.qwen_2_5 import Qwen2_5
from models.deepseek_v3 import DeepSeek_V3
from models.tiny_aya import TinyAya

# model defination
# copied from /models


def build_models(hf_token: str) -> dict[str, Any]:
    return {
        "Llama-3-8B": Llama3_8B(token=hf_token),
        "Mistral-7B": Mistral_7b(token=hf_token),
        "Qwen-2.5": Qwen2_5(token=hf_token),
        "DeepSeek-V3": DeepSeek_V3(token=hf_token),
        "TinyAya": TinyAya(token=hf_token),
    }


def resolve_model(name: str, models: dict[str, Any]) -> tuple[str, Any]:
    aliases = {
        "llama": "Llama-3-8B",
        "mistral": "Mistral-7B",
        "qwen": "Qwen-2.5",
        "deepseek": "DeepSeek-V3",
        "tinyaya": "TinyAya",
    }
    model_key = aliases.get(name.lower(), name)
    if model_key not in models:
        allowed = ", ".join(models.keys())
        raise HTTPException(status_code=400, detail=f"Unknown model '{name}'. Use one of: {allowed}")
    return model_key, models[model_key]