llm.py

from smolagents import LiteLLMModel, TransformersModel


ANTHROPIC_MODEL_IDS = [
    "claude-opus-4-20250514",
    "claude-sonnet-4-20250514",
    "claude-3-7-sonnet-latest",
    "claude-3-5-haiku-latest",
    "claude-3-5-sonnet-latest",
    "claude-3-5-sonnet-20240620",
    "claude-3-opus-latest",
    "claude-3-sonnet-20240229",
    "claude-3-haiku-20240307",
]

TRANSFORMERS_MODEL_IDS = [
    "Qwen/Qwen2.5-VL-3B-Instruct",
    "Qwen/Qwen2.5-VL-7B-Instruct",
    "Qwen/Qwen2.5-VL-14B-Instruct",
    "Qwen/Qwen2.5-VL-32B-Instruct",
    "Qwen/Qwen2.5-VL-72B-Instruct",
]


def get_anthropic_model(model_id, anthropic_api_key):
    if model_id not in ANTHROPIC_MODEL_IDS:
        raise ValueError(f"Model {model_id} not found in Anthropic model IDs")
    model = LiteLLMModel(
        model_id=model_id,
        api_key=anthropic_api_key,
    )
    return model


def get_transformers_model(model_id):
    model = TransformersModel(model_id=model_id)
    return model