tests/test_llms.py

import pytest

from gaia_agent.config import Settings
from gaia_agent.llms import create_chat_model


def test_create_litellm_chat_model():
    model = create_chat_model(
        Settings(
            litellm_model="anthropic/claude-3-5-sonnet-latest",
            litellm_api_key="test-key",
        )
    )

    assert type(model).__name__ == "ChatLiteLLM"
    assert model.model == "anthropic/claude-3-5-sonnet-latest"


def test_create_litellm_chat_model_supports_proxy_settings():
    model = create_chat_model(
        Settings(
            litellm_model="gaia-router",
            litellm_api_key="test-key",
            litellm_api_base="http://localhost:4000",
        )
    )

    assert type(model).__name__ == "ChatLiteLLM"
    assert model.model == "gaia-router"
    assert model.api_base == "http://localhost:4000"


def test_create_chat_model_requires_litellm_model():
    with pytest.raises(ValueError, match="LITELLM_MODEL must be set"):
        create_chat_model(Settings(litellm_model=None))