gaia_agent/llms.py

from langchain_core.language_models.chat_models import BaseChatModel
from langchain_litellm import ChatLiteLLM

from gaia_agent.config import Settings, settings


def create_chat_model(config: Settings = settings) -> BaseChatModel:
    """Create the configured LiteLLM chat model."""
    if not config.litellm_model:
        raise ValueError("LITELLM_MODEL must be set to create a chat model.")

    kwargs = {
        "model": config.litellm_model,
        "temperature": config.litellm_temperature,
    }
    if config.litellm_api_key:
        kwargs["api_key"] = config.litellm_api_key
    if config.litellm_api_base:
        kwargs["api_base"] = config.litellm_api_base

    return ChatLiteLLM(**kwargs)