src/generate/llms.py

from typing import Literal

from langchain_ollama import ChatOllama
from langchain_openai.chat_models import ChatOpenAI
from langchain_community.chat_models import GigaChat
from langchain_core.language_models.chat_models import BaseChatModel


LLMName = Literal["ollama", "openai", "gigachat"]
LLM_NAME_TO_CLASS: dict[LLMName, type[BaseChatModel]] = {
    "ollama": ChatOllama,
    "openai": ChatOpenAI,
    "gigachat": GigaChat,
}