models.py

import os
from smolagents import LiteLLMModel
from dotenv import load_dotenv

load_dotenv()


def select_model(provider, name_model=None):
    match provider:
        case "mistral":
            reference_model = "mistral-small-latest"
            api_key = os.getenv("MISTRAL_API_KEY")
        case "openrouter":
            match name_model:
                case "gemini":
                    reference_model = "google/gemini-2.5-pro-exp-03-25"
                case "deepseek":
                    reference_model = "deepseek/deepseek-chat-v3-0324:free"
                case _:
                    print("Invalid command.")
            api_key = os.getenv("OPENROUTER_API_KEY")
        case "groq":
            match name_model:
                case "llama":
                    reference_model = "llama-3.3-70b-versatile"
                case "deepseek":
                    reference_model = "deepseek-r1-distill-llama-70b"
                case _:
                    print("Invalid command.")
            api_key = os.getenv("GROQ_API_KEY")
        case _:
            print("Invalid command.")

    return LiteLLMModel(
        model_id=f"{provider}/{reference_model}", api_key=api_key
    )