import os
import warnings
from typing import Literal, Optional

from pydantic import BaseModel, ConfigDict, model_validator


class Model(BaseModel):
    """
    Model information.

    Attributes:
        id: Unique identifier for the model, passed as model parameter for requests
        model_type: API type
        client: Client name
        endpoint: Custom endpoint for the model
        aliases: List of aliases for the model
        supports_tools: Whether the model supports tool calling
        supports_structured_output: Whether the model supports structured output
        supports_thinking: Whether the model supports configuring thinking
            on/off through system message or request parameters
        thinking_prefix: System message prefix when thinking is enabled
            (tag-based)
        no_thinking_prefix: System message prefix when thinking is disabled
            (tag-based)
        thinking_param_enable: Dict of parameters to apply when thinking is
            enabled (param-based)
        thinking_param_disable: Dict of parameters to apply when thinking is
            disabled (param-based)

    Thinking mode can be enabled via two mechanisms:
        1. Tag-based: Use thinking_prefix/no_thinking_prefix (appended to
            system message)
        2. Param-based: Use thinking_param_enable/thinking_param_disable
            (merged into request params)

    All aliases are deprecated and will trigger a warning when used.
    """

    # https://docs.pydantic.dev/latest/api/config/#protected-namespaces
    # we disable this protection to avoid warnings from `model_type`
    model_config = ConfigDict(
        protected_namespaces=(),
    )

    id: str

    # why do we have a model_type? because ChatNVIDIA can speak both chat and vlm.
    model_type: Optional[
        Literal["chat", "vlm", "nv-vlm", "embedding", "ranking", "completions", "qa"]
    ] = None

    client: Optional[
        Literal["ChatNVIDIA", "NVIDIAEmbeddings", "NVIDIARerank", "NVIDIA"]
    ] = None

    endpoint: Optional[str] = None

    aliases: Optional[list] = None

    supports_tools: Optional[bool] = False

    supports_structured_output: Optional[bool] = False

    supports_thinking: Optional[bool] = False

    thinking_prefix: Optional[str] = None

    no_thinking_prefix: Optional[str] = None

    thinking_param_enable: Optional[dict] = None

    thinking_param_disable: Optional[dict] = None

    base_model: Optional[str] = None

    def __hash__(self) -> int:
        return hash(self.id)

    @model_validator(mode="after")
    def validate_client(self) -> "Model":
        if self.client:
            supported = {
                "ChatNVIDIA": ("chat", "vlm", "nv-vlm", "qa"),
                "NVIDIAEmbeddings": ("embedding",),
                "NVIDIARerank": ("ranking",),
                "NVIDIA": ("completions",),
            }
            if self.model_type not in supported.get(self.client, ()):
                raise ValueError(
                    f"Model type '{self.model_type}' not supported "
                    f"by client '{self.client}'"
                )
        return self

    @model_validator(mode="after")
    def validate_thinking_config(self) -> "Model":
        """Warn if both param-based and tag-based thinking are configured."""
        has_param_based = (
            self.thinking_param_enable is not None
            or self.thinking_param_disable is not None
        )
        has_tag_based = (
            self.thinking_prefix is not None or self.no_thinking_prefix is not None
        )

        if has_param_based and has_tag_based:
            warnings.warn(
                f"Model '{self.id}' has both param-based thinking "
                f"(thinking_param_enable/disable) and tag-based thinking "
                f"(thinking_prefix/no_thinking_prefix) configured. "
                f"Param-based thinking will take precedence and tag-based "
                f"thinking will be ignored.",
                UserWarning,
                stacklevel=2,
            )
        return self


CHAT_MODEL_TABLE = {
    "google/gemma-7b": Model(
        id="google/gemma-7b",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-gemma-7b", "playground_gemma_7b", "gemma_7b"],
    ),
    "mistralai/mistral-7b-instruct-v0.2": Model(
        id="mistralai/mistral-7b-instruct-v0.2",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-mistral-7b-instruct-v2", "playground_mistral_7b", "mistral_7b"],
    ),
    "mistralai/mixtral-8x7b-instruct-v0.1": Model(
        id="mistralai/mixtral-8x7b-instruct-v0.1",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-mixtral-8x7b-instruct", "playground_mixtral_8x7b", "mixtral_8x7b"],
    ),
    "google/codegemma-7b": Model(
        id="google/codegemma-7b",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-codegemma-7b"],
    ),
    "gotocompany/gemma-2-9b-cpt-sahabatai-instruct": Model(
        id="gotocompany/gemma-2-9b-cpt-sahabatai-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "mistralai/mixtral-8x22b-instruct-v0.1": Model(
        id="mistralai/mixtral-8x22b-instruct-v0.1",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-mixtral-8x22b-instruct"],
    ),
    "mistralai/mixtral-8x22b-v0.1": Model(
        id="mistralai/mixtral-8x22b-v0.1",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "meta/llama3-8b-instruct": Model(
        id="meta/llama3-8b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-llama3-8b"],
    ),
    "meta/llama3-70b-instruct": Model(
        id="meta/llama3-70b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-llama3-70b"],
    ),
    "microsoft/phi-3-mini-128k-instruct": Model(
        id="microsoft/phi-3-mini-128k-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-phi-3-mini"],
    ),
    "microsoft/phi-3-mini-4k-instruct": Model(
        id="microsoft/phi-3-mini-4k-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-phi-3-mini-4k", "playground_phi2", "phi2"],
    ),
    "aisingapore/sea-lion-7b-instruct": Model(
        id="aisingapore/sea-lion-7b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-sea-lion-7b-instruct"],
    ),
    "microsoft/phi-3-small-8k-instruct": Model(
        id="microsoft/phi-3-small-8k-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-phi-3-small-8k-instruct"],
    ),
    "microsoft/phi-3-small-128k-instruct": Model(
        id="microsoft/phi-3-small-128k-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-phi-3-small-128k-instruct"],
    ),
    "microsoft/phi-3-medium-4k-instruct": Model(
        id="microsoft/phi-3-medium-4k-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-phi-3-medium-4k-instruct"],
    ),
    "ibm/granite-3.3-8b-instruct": Model(
        id="ibm/granite-3.3-8b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_thinking=True,
        supports_tools=True,
        thinking_param_enable={"chat_template_kwargs": {"enable_thinking": True}},
        thinking_param_disable={"chat_template_kwargs": {"enable_thinking": False}},
    ),
    "ibm/granite-guardian-3.0-8b": Model(
        id="ibm/granite-guardian-3.0-8b",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "mediatek/breeze-7b-instruct": Model(
        id="mediatek/breeze-7b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-breeze-7b-instruct"],
    ),
    "upstage/solar-10.7b-instruct": Model(
        id="upstage/solar-10.7b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-solar-10_7b-instruct"],
    ),
    "mistralai/mistral-7b-instruct-v0.3": Model(
        id="mistralai/mistral-7b-instruct-v0.3",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-mistral-7b-instruct-v03"],
    ),
    "01-ai/yi-large": Model(
        id="01-ai/yi-large",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-yi-large"],
    ),
    "mistralai/codestral-22b-instruct-v0.1": Model(
        id="mistralai/codestral-22b-instruct-v0.1",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-codestral-22b-instruct-v01"],
        supports_structured_output=True,
    ),
    "google/gemma-2-9b-it": Model(
        id="google/gemma-2-9b-it",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-gemma-2-9b-it"],
    ),
    "google/gemma-2-27b-it": Model(
        id="google/gemma-2-27b-it",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-gemma-2-27b-it"],
    ),
    "google/gemma-3-12b-it": Model(
        id="google/gemma-3-12b-it",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "google/gemma-3-1b-it": Model(
        id="google/gemma-3-1b-it",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "google/gemma-3-27b-it": Model(
        id="google/gemma-3-27b-it",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "google/gemma-3-4b-it": Model(
        id="google/gemma-3-4b-it",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "google/gemma-3n-e2b-it": Model(
        id="google/gemma-3n-e2b-it",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
    "google/gemma-3n-e4b-it": Model(
        id="google/gemma-3n-e4b-it",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
    "google/shieldgemma-9b": Model(
        id="google/shieldgemma-9b",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "microsoft/phi-3-medium-128k-instruct": Model(
        id="microsoft/phi-3-medium-128k-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        aliases=["ai-phi-3-medium-128k-instruct"],
    ),
    "nv-mistralai/mistral-nemo-12b-instruct": Model(
        id="nv-mistralai/mistral-nemo-12b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "meta/llama-3.1-8b-instruct": Model(
        id="meta/llama-3.1-8b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "meta/llama-3.1-70b-instruct": Model(
        id="meta/llama-3.1-70b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "meta/llama-3.1-405b-instruct": Model(
        id="meta/llama-3.1-405b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "meta/llama-guard-4-12b": Model(
        id="meta/llama-guard-4-12b",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "mistralai/mamba-codestral-7b-v0.1": Model(
        id="mistralai/mamba-codestral-7b-v0.1",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "google/gemma-2-2b-it": Model(
        id="google/gemma-2-2b-it",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "mistralai/magistral-small-2506": Model(
        id="mistralai/magistral-small-2506",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "mistralai/mistral-medium-3-instruct": Model(
        id="mistralai/mistral-medium-3-instruct",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
    "mistralai/mistral-nemotron": Model(
        id="mistralai/mistral-nemotron",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
    ),
    "mistralai/mistral-small-24b-instruct": Model(
        id="mistralai/mistral-small-24b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "mistralai/mistral-small-3.1-24b-instruct-2503": Model(
        id="mistralai/mistral-small-3.1-24b-instruct-2503",
        model_type="vlm",
        client="ChatNVIDIA",
        supports_tools=True,
    ),
    "mistralai/mathstral-7b-v0.1": Model(
        id="mistralai/mathstral-7b-v0.1",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "rakuten/rakutenai-7b-instruct": Model(
        id="rakuten/rakutenai-7b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "rakuten/rakutenai-7b-chat": Model(
        id="rakuten/rakutenai-7b-chat",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "baichuan-inc/baichuan2-13b-chat": Model(
        id="baichuan-inc/baichuan2-13b-chat",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "thudm/chatglm3-6b": Model(
        id="thudm/chatglm3-6b",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "tiiuae/falcon3-7b-instruct": Model(
        id="tiiuae/falcon3-7b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "microsoft/phi-3.5-mini-instruct": Model(
        id="microsoft/phi-3.5-mini-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "microsoft/phi-4-mini-flash-reasoning": Model(
        id="microsoft/phi-4-mini-flash-reasoning",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "microsoft/phi-4-mini-instruct": Model(
        id="microsoft/phi-4-mini-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
    ),
    "microsoft/phi-4-multimodal-instruct": Model(
        id="microsoft/phi-4-multimodal-instruct",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
    "utter-project/eurollm-9b-instruct": Model(
        id="utter-project/eurollm-9b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "nvidia/nemotron-mini-4b-instruct": Model(
        id="nvidia/nemotron-mini-4b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "yentinglin/llama-3-taiwan-70b-instruct": Model(
        id="yentinglin/llama-3-taiwan-70b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "tokyotech-llm/llama-3-swallow-70b-instruct-v0.1": Model(
        id="tokyotech-llm/llama-3-swallow-70b-instruct-v0.1",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "abacusai/dracarys-llama-3.1-70b-instruct": Model(
        id="abacusai/dracarys-llama-3.1-70b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "qwen/qwen2-7b-instruct": Model(
        id="qwen/qwen2-7b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "meta/llama-3.2-1b-instruct": Model(
        id="meta/llama-3.2-1b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "meta/llama-3.2-3b-instruct": Model(
        id="meta/llama-3.2-3b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "institute-of-science-tokyo/llama-3.1-swallow-8b-instruct-v0.1": Model(
        id="institute-of-science-tokyo/llama-3.1-swallow-8b-instruct-v0.1",
        model_type="chat",
        client="ChatNVIDIA",
        supports_structured_output=True,
    ),
    "institute-of-science-tokyo/llama-3.1-swallow-70b-instruct-v0.1": Model(
        id="institute-of-science-tokyo/llama-3.1-swallow-70b-instruct-v0.1",
        model_type="chat",
        client="ChatNVIDIA",
        supports_structured_output=True,
    ),
    "igenius/colosseum_355b_instruct_16k": Model(
        id="igenius/colosseum_355b_instruct_16k",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "igenius/italia_10b_instruct_16k": Model(
        id="igenius/italia_10b_instruct_16k",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "zyphra/zamba2-7b-instruct": Model(
        id="zyphra/zamba2-7b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "nvidia/nemotron-4-mini-hindi-4b-instruct": Model(
        id="nvidia/nemotron-4-mini-hindi-4b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_structured_output=True,
    ),
    "nvidia/llama-3.1-nemoguard-8b-content-safety": Model(
        id="nvidia/llama-3.1-nemoguard-8b-content-safety",
        model_type="chat",
        client="ChatNVIDIA",
        supports_structured_output=True,
    ),
    "nvidia/llama-3.1-nemoguard-8b-topic-control": Model(
        id="nvidia/llama-3.1-nemoguard-8b-topic-control",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "nvidia/usdcode-llama-3.1-70b-instruct": Model(
        id="nvidia/usdcode-llama-3.1-70b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "meta/llama-3.3-70b-instruct": Model(
        id="meta/llama-3.3-70b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "qwen/qwen2.5-coder-32b-instruct": Model(
        id="qwen/qwen2.5-coder-32b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "qwen/qwen2.5-coder-7b-instruct": Model(
        id="qwen/qwen2.5-coder-7b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "qwen/qwen2.5-7b-instruct": Model(
        id="qwen/qwen2.5-7b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "qwen/qwen3-235b-a22b": Model(
        id="qwen/qwen3-235b-a22b",
        model_type="chat",
        client="ChatNVIDIA",
        supports_thinking=True,
        supports_tools=True,
        thinking_param_enable={"chat_template_kwargs": {"enable_thinking": True}},
        thinking_param_disable={"chat_template_kwargs": {"enable_thinking": False}},
    ),
    "qwen/qwq-32b": Model(
        id="qwen/qwq-32b",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
    ),
    "nvidia/llama-3.1-nemotron-70b-reward": Model(
        id="nvidia/llama-3.1-nemotron-70b-reward",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "deepseek-ai/deepseek-r1": Model(
        id="deepseek-ai/deepseek-r1",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "deepseek-ai/deepseek-r1-0528": Model(
        id="deepseek-ai/deepseek-r1-0528",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
    ),
    "deepseek-ai/deepseek-r1-distill-llama-8b": Model(
        id="deepseek-ai/deepseek-r1-distill-llama-8b",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "deepseek-ai/deepseek-r1-distill-qwen-14b": Model(
        id="deepseek-ai/deepseek-r1-distill-qwen-14b",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "deepseek-ai/deepseek-r1-distill-qwen-32b": Model(
        id="deepseek-ai/deepseek-r1-distill-qwen-32b",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "deepseek-ai/deepseek-r1-distill-qwen-7b": Model(
        id="deepseek-ai/deepseek-r1-distill-qwen-7b",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "nvidia/llama-3.1-nemotron-nano-8b-v1": Model(
        id="nvidia/llama-3.1-nemotron-nano-8b-v1",
        model_type="chat",
        client="ChatNVIDIA",
        supports_thinking=True,
        thinking_prefix="detailed thinking on",
        no_thinking_prefix="detailed thinking off",
    ),
    "nvidia/llama-3.1-nemotron-nano-4b-v1.1": Model(
        id="nvidia/llama-3.1-nemotron-nano-4b-v1.1",
        model_type="chat",
        client="ChatNVIDIA",
        supports_thinking=True,
        thinking_prefix="detailed thinking on",
        no_thinking_prefix="detailed thinking off",
        supports_tools=True,
    ),
    "nvidia/llama-3.1-nemotron-ultra-253b-v1": Model(
        id="nvidia/llama-3.1-nemotron-ultra-253b-v1",
        model_type="chat",
        client="ChatNVIDIA",
        supports_thinking=True,
        thinking_prefix="detailed thinking on",
        no_thinking_prefix="detailed thinking off",
        supports_tools=True,
    ),
    "nvidia/llama-3.3-nemotron-super-49b-v1": Model(
        id="nvidia/llama-3.3-nemotron-super-49b-v1",
        model_type="chat",
        client="ChatNVIDIA",
        supports_thinking=True,
        thinking_prefix="detailed thinking on",
        no_thinking_prefix="detailed thinking off",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "nvidia/llama-3.3-nemotron-super-49b-v1.5": Model(
        id="nvidia/llama-3.3-nemotron-super-49b-v1.5",
        model_type="chat",
        client="ChatNVIDIA",
        supports_thinking=True,
        thinking_prefix="/think",
        no_thinking_prefix="/no_think",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "moonshotai/kimi-k2-instruct": Model(
        id="moonshotai/kimi-k2-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "openai/gpt-oss-20b": Model(
        id="openai/gpt-oss-20b",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "openai/gpt-oss-120b": Model(
        id="openai/gpt-oss-120b",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "marin/marin-8b-instruct": Model(
        id="marin/marin-8b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "speakleash/bielik-11b-v2.3-instruct": Model(
        id="speakleash/bielik-11b-v2.3-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "nvidia/nvidia-nemotron-nano-9b-v2": Model(
        id="nvidia/nvidia-nemotron-nano-9b-v2",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
        supports_thinking=True,
        thinking_prefix="/think",
        no_thinking_prefix="/no_think",
    ),
    "deepseek-ai/deepseek-v3.1": Model(
        id="deepseek-ai/deepseek-v3.1",
        model_type="chat",
        client="ChatNVIDIA",
        supports_thinking=True,
        thinking_param_enable={"chat_template_kwargs": {"enable_thinking": True}},
        thinking_param_disable={"chat_template_kwargs": {"enable_thinking": False}},
    ),
    "bytedance/seed-oss-36b-instruct": Model(
        id="bytedance/seed-oss-36b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "moonshotai/kimi-k2-instruct-0905": Model(
        id="moonshotai/kimi-k2-instruct-0905",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "qwen/qwen3-next-80b-a3b-instruct": Model(
        id="qwen/qwen3-next-80b-a3b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "qwen/qwen3-next-80b-a3b-thinking": Model(
        id="qwen/qwen3-next-80b-a3b-thinking",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "deepseek-ai/deepseek-v3.1-terminus": Model(
        id="deepseek-ai/deepseek-v3.1-terminus",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
        supports_thinking=True,
        thinking_param_enable={"chat_template_kwargs": {"enable_thinking": True}},
        thinking_param_disable={"chat_template_kwargs": {"enable_thinking": False}},
    ),
    "minimaxai/minimax-m2": Model(
        id="minimaxai/minimax-m2",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "moonshotai/kimi-k2-thinking": Model(
        id="moonshotai/kimi-k2-thinking",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
    ),
    "nvidia/nemotron-3-nano-30b-a3b": Model(
        id="nvidia/nemotron-3-nano-30b-a3b",
        model_type="chat",
        client="ChatNVIDIA",
        supports_tools=True,
        supports_structured_output=True,
        supports_thinking=True,
        thinking_param_enable={"chat_template_kwargs": {"enable_thinking": True}},
        thinking_param_disable={"chat_template_kwargs": {"enable_thinking": False}},
    ),
    "deepseek-ai/deepseek-v3.2": Model(
        id="deepseek-ai/deepseek-v3.2",
        model_type="chat",
        client="ChatNVIDIA",
        supports_thinking=True,
        supports_structured_output=True,
        supports_tools=True,
        thinking_param_enable={"chat_template_kwargs": {"enable_thinking": True}},
        thinking_param_disable={"chat_template_kwargs": {"enable_thinking": False}},
    ),
}

QA_MODEL_TABLE = {
    "nvidia/llama3-chatqa-1.5-8b": Model(
        id="nvidia/llama3-chatqa-1.5-8b",
        model_type="qa",
        client="ChatNVIDIA",
        aliases=["ai-chatqa-1.5-8b"],
    ),
}

VLM_MODEL_TABLE = {
    "google/paligemma": Model(
        id="google/paligemma",
        model_type="nv-vlm",
        client="ChatNVIDIA",
        endpoint="https://ai.api.nvidia.com/v1/vlm/google/paligemma",
        aliases=["ai-google-paligemma"],
    ),
    "microsoft/phi-3-vision-128k-instruct": Model(
        id="microsoft/phi-3-vision-128k-instruct",
        model_type="nv-vlm",
        client="ChatNVIDIA",
        endpoint="https://ai.api.nvidia.com/v1/vlm/microsoft/phi-3-vision-128k-instruct",
        aliases=["ai-phi-3-vision-128k-instruct"],
    ),
    "microsoft/phi-3.5-vision-instruct": Model(
        id="microsoft/phi-3.5-vision-instruct",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
    "nvidia/nvclip": Model(
        id="nvidia/nvclip",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
    "nvidia/nemoretriever-parse": Model(
        id="nvidia/nemoretriever-parse",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "nvidia/riva-translate-4b-instruct": Model(
        id="nvidia/riva-translate-4b-instruct",
        model_type="chat",
        client="ChatNVIDIA",
    ),
    "meta/llama-3.2-11b-vision-instruct": Model(
        id="meta/llama-3.2-11b-vision-instruct",
        model_type="vlm",
        client="ChatNVIDIA",
        endpoint="https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-11b-vision-instruct/chat/completions",
    ),
    "meta/llama-3.2-90b-vision-instruct": Model(
        id="meta/llama-3.2-90b-vision-instruct",
        model_type="vlm",
        client="ChatNVIDIA",
        endpoint="https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct/chat/completions",
    ),
    "meta/llama-4-maverick-17b-128e-instruct": Model(
        id="meta/llama-4-maverick-17b-128e-instruct",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
    "meta/llama-4-scout-17b-16e-instruct": Model(
        id="meta/llama-4-scout-17b-16e-instruct",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
    "nvidia/llama-3.1-nemotron-nano-vl-8b-v1": Model(
        id="nvidia/llama-3.1-nemotron-nano-vl-8b-v1",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
    "mistralai/ministral-14b-instruct-2512": Model(
        id="mistralai/ministral-14b-instruct-2512",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
    "mistralai/mistral-large-3-675b-instruct-2512": Model(
        id="mistralai/mistral-large-3-675b-instruct-2512",
        model_type="vlm",
        client="ChatNVIDIA",
    ),
}

EMBEDDING_MODEL_TABLE = {
    "snowflake/arctic-embed-l": Model(
        id="snowflake/arctic-embed-l",
        model_type="embedding",
        client="NVIDIAEmbeddings",
        aliases=["ai-arctic-embed-l"],
    ),
    "NV-Embed-QA": Model(
        id="NV-Embed-QA",
        model_type="embedding",
        client="NVIDIAEmbeddings",
        endpoint="https://ai.api.nvidia.com/v1/retrieval/nvidia/embeddings",
        aliases=[
            "ai-embed-qa-4",
        ],
    ),
    "nvidia/nv-embed-v1": Model(
        id="nvidia/nv-embed-v1",
        model_type="embedding",
        client="NVIDIAEmbeddings",
        aliases=["ai-nv-embed-v1"],
    ),
    "nvidia/nv-embedqa-mistral-7b-v2": Model(
        id="nvidia/nv-embedqa-mistral-7b-v2",
        model_type="embedding",
        client="NVIDIAEmbeddings",
    ),
    "nvidia/nv-embedqa-e5-v5": Model(
        id="nvidia/nv-embedqa-e5-v5",
        model_type="embedding",
        client="NVIDIAEmbeddings",
    ),
    "baai/bge-m3": Model(
        id="baai/bge-m3",
        model_type="embedding",
        client="NVIDIAEmbeddings",
    ),
    "nvidia/llama-3.2-nv-embedqa-1b-v1": Model(
        id="nvidia/llama-3.2-nv-embedqa-1b-v1",
        model_type="embedding",
        client="NVIDIAEmbeddings",
    ),
    "nvidia/llama-3.2-nv-embedqa-1b-v2": Model(
        id="nvidia/llama-3.2-nv-embedqa-1b-v2",
        model_type="embedding",
        client="NVIDIAEmbeddings",
    ),
    "nvidia/llama-3.2-nemoretriever-1b-vlm-embed-v1": Model(
        id="nvidia/llama-3.2-nemoretriever-1b-vlm-embed-v1",
        model_type="embedding",
        client="NVIDIAEmbeddings",
    ),
    "nvidia/nv-embedcode-7b-v1": Model(
        id="nvidia/nv-embedcode-7b-v1",
        model_type="embedding",
        client="NVIDIAEmbeddings",
    ),
    "nvidia/llama-3.2-nemoretriever-300m-embed-v2": Model(
        id="nvidia/llama-3.2-nemoretriever-300m-embed-v2",
        model_type="embedding",
        client="NVIDIAEmbeddings",
    ),
}

RANKING_MODEL_TABLE = {
    "nv-rerank-qa-mistral-4b:1": Model(
        id="nv-rerank-qa-mistral-4b:1",
        model_type="ranking",
        client="NVIDIARerank",
        endpoint="https://ai.api.nvidia.com/v1/retrieval/nvidia/reranking",
        aliases=["ai-rerank-qa-mistral-4b"],
    ),
    "nvidia/llama-3.2-nv-rerankqa-1b-v1": Model(
        id="nvidia/llama-3.2-nv-rerankqa-1b-v1",
        model_type="ranking",
        client="NVIDIARerank",
        endpoint="https://ai.api.nvidia.com/v1/retrieval/nvidia/llama-3_2-nv-rerankqa-1b-v1/reranking",
    ),
    "nvidia/llama-3.2-nv-rerankqa-1b-v2": Model(
        id="nvidia/llama-3.2-nv-rerankqa-1b-v2",
        model_type="ranking",
        client="NVIDIARerank",
        endpoint="https://ai.api.nvidia.com/v1/retrieval/nvidia/llama-3_2-nv-rerankqa-1b-v2/reranking",
    ),
}

COMPLETION_MODEL_TABLE = {
    "bigcode/starcoder2-7b": Model(
        id="bigcode/starcoder2-7b",
        model_type="completions",
        client="NVIDIA",
    ),
    "nvidia/mistral-nemo-minitron-8b-base": Model(
        id="nvidia/mistral-nemo-minitron-8b-base",
        model_type="completions",
        client="NVIDIA",
    ),
}


OPENAI_MODEL_TABLE = {
    "gpt-3.5-turbo": Model(
        id="gpt-3.5-turbo",
        model_type="chat",
        client="ChatNVIDIA",
        endpoint="https://api.openai.com/v1/chat/completions",
        supports_tools=True,
    ),
}


MODEL_TABLE = {
    **CHAT_MODEL_TABLE,
    **QA_MODEL_TABLE,
    **VLM_MODEL_TABLE,
    **EMBEDDING_MODEL_TABLE,
    **RANKING_MODEL_TABLE,
    **COMPLETION_MODEL_TABLE,
}

if "_INCLUDE_OPENAI" in os.environ:
    MODEL_TABLE.update(OPENAI_MODEL_TABLE)


def register_model(model: Model) -> None:
    """Register a model as a known model.

    Must be done at the beginning of a program, at least before the model is used or
    available models are listed.

    For instance:

    ```python
    from langchain_nvidia_ai_endpoints import register_model, Model


    register_model(
        Model(
            id="my-custom-model-name",
            model_type="chat",
            client="ChatNVIDIA",
            endpoint="http://host:port/path-to-my-model"
        )
    )
    llm = ChatNVIDIA(model="my-custom-model-name")
    ```

    Be sure that the `id` matches the model parameter the endpoint expects.

    Supported model types are chat models, which must accept and produce chat completion
    payloads.

    Supported model clients are `ChatNVIDIA`, for chat models.

    Endpoint is required.

    Use this instead of passing `base_url` to a client constructor when the model's
    endpoint supports inference and not `/v1/models` listing.

    Use `base_url` when the model's endpoint supports `/v1/models` listing and inference
    on a known path, e.g. `/v1/chat/completions`.
    """
    if model.id in MODEL_TABLE:
        warnings.warn(
            f"Model {model.id} is already registered. "
            f"Overriding {MODEL_TABLE[model.id]}",
            UserWarning,
        )
    if not model.endpoint:
        raise ValueError(f"Model {model.id} does not have an endpoint.")
    MODEL_TABLE[model.id] = model


def lookup_model(name: str) -> Optional[Model]:
    """Lookup a model by name, using only the table of known models.

    The name is either:

    - Directly in the table
    - An alias in the table
    - Not found (`None`)

    Callers can check to see if the name was an alias by comparing the result's id field
    to the name they provided.
    """
    model = None
    if not (model := MODEL_TABLE.get(name)):
        for mdl in MODEL_TABLE.values():
            if mdl.aliases and name in mdl.aliases:
                model = mdl
                break
    return model


def determine_model(name: str) -> Optional[Model]:
    """Determine the model to use based on a name, using only the table of known models.

    Raise a warning if the model is found to be an alias of a known model.

    If the model is not found, return `None`.
    """
    if model := lookup_model(name):
        # all aliases are deprecated
        if model.id != name:
            warnings.warn(
                f"Model {name} is deprecated. Using {model.id} instead.", UserWarning
            )
    return model
