litellm/tests/unit/test_utils_get_model_info.py

import json
import re
from collections.abc import Collection, Iterator, Mapping
from pathlib import Path
from typing import Final, Literal, cast

import httpx
import pytest
import respx

import litellm
from litellm import get_model_info
from litellm.llms.bedrock.common_utils import BedrockModelInfo
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.types.utils import ModelInfoBase
from litellm.utils import _invalidate_model_cost_lowercase_map
import os
from unittest.mock import MagicMock, patch


@pytest.fixture(autouse=True)
def isolate_model_info_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
    monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
    monkeypatch.setattr(litellm, "custom_provider_map", list(litellm.custom_provider_map))
    yield
    _invalidate_model_cost_lowercase_map()


def test_get_model_info_simple_model_name():
    """
    tests if model name given, and model exists in model info - the object is returned
    """
    model = "claude-opus-5-5"
    litellm.get_model_info(model)


def test_get_model_info_custom_llm_with_model_name():
    """
    Tests if {custom_llm_provider}/{model_name} name given, and model exists in model info, the object is returned
    """
    model = "anthropic/claude-opus-5-5"
    litellm.get_model_info(model)


def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch):
    """
    Tests if {custom_llm_provider}/{model_name} name given, and model exists in model info, the object is returned
    """
    model = "command-r-plus"
    provider = "openai"  # vllm is openai-compatible
    litellm.register_model(
        {
            "openai/command-r-plus": {
                "input_cost_per_token": 0.0,
                "output_cost_per_token": 0.0,
            },
        }
    )
    model_info = litellm.get_model_info(model, custom_llm_provider=provider)
    print("model_info", model_info)
    assert model_info["input_cost_per_token"] == 0.0


def test_get_model_info_ollama_chat():
    from litellm.llms.ollama.completion.transformation import OllamaConfig

    with patch.object(
        litellm.module_level_client,
        "post",
        return_value=MagicMock(
            json=lambda: {
                "model_info": {"llama.context_length": 32768},
                "template": "tools",
            }
        ),
    ) as mock_client:
        info = OllamaConfig().get_model_info("unknown-model")
        assert info["supports_function_calling"] is True

        info = get_model_info("ollama/unknown-model")
        print("info", info)
        assert info["supports_function_calling"] is True

        mock_client.assert_called()

        print(mock_client.call_args.kwargs)

        assert mock_client.call_args.kwargs["json"]["name"] == "unknown-model"


def test_get_model_info_bedrock_region(monkeypatch):
    regional_model = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
    model_cost_without_regional_entry = {
        key: value for key, value in litellm.get_model_cost_map(url="").items() if key != regional_model
    }
    monkeypatch.setattr(litellm, "model_cost", model_cost_without_regional_entry)
    _invalidate_model_cost_lowercase_map()
    info = litellm.get_model_info(model=regional_model, custom_llm_provider="bedrock")
    print("info", info)
    assert info["key"] == "anthropic.claude-haiku-4-5-20251001-v1:0"
    assert info["litellm_provider"] == "bedrock_converse"


@pytest.mark.parametrize(
    "model",
    [
        "ft:gpt-3.5-turbo:my-org:custom_suffix:id",
        "ft:gpt-4-0613:my-org:custom_suffix:id",
        "ft:davinci-002:my-org:custom_suffix:id",
        "ft:babbage-002:my-org:custom_suffix:id",
        "gpt-35-turbo",
        "ada",
    ],
)
def test_get_model_info_completion_cost_unit_tests(model):
    info = litellm.get_model_info(model)
    print("info", info)


def test_get_model_info_ft_model_with_provider_prefix():
    args = {
        "model": "openai/ft:gpt-3.5-turbo:my-org:custom_suffix:id",
        "custom_llm_provider": "openai",
    }
    info = litellm.get_model_info(**args)
    print("info", info)
    assert info["key"] == "ft:gpt-3.5-turbo"


def _enforce_bedrock_converse_models(
    model_cost: Mapping[str, ModelInfoBase], whitelist_models: Collection[str]
) -> None:
    for model, info in model_cost.items():
        if (
            info.get("litellm_provider") == "bedrock"
            and info.get("mode") == "chat"
            and model not in whitelist_models
            and not (
                (base_model := BedrockModelInfo.get_base_model(model)) != model
                and model_cost.get(base_model, {}).get("litellm_provider") == "bedrock_converse"
                and BedrockModelInfo.get_bedrock_route(model) == "converse"
            )
        ):
            raise AssertionError(f"Unlisted Bedrock chat model does not route to Converse: {model}")


def _read_whitelisted_bedrock_models() -> tuple[str, ...]:
    path: Final = Path(__file__).resolve().parents[2] / "whitelisted_bedrock_models.txt"
    return tuple(path.read_text().splitlines())


def _normalize_bedrock_model_key(model_key: str) -> str:
    without_wildcard: Final = model_key.replace("*/", "")
    return re.sub(r"(?:1-month-commitment|3-month-commitment|6-month-commitment)/", "", without_wildcard)


def test_model_info_bedrock_converse(monkeypatch):
    """
    Assert unlisted Bedrock chat models declare or inherit Converse routing.

    This ensures they are automatically routed to the converse endpoint.
    """
    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
    litellm.model_cost = litellm.get_model_cost_map(url="")
    try:
        # Load whitelist models from file
        with open("whitelisted_bedrock_models.txt", "r") as file:
            whitelist_models = [line.strip() for line in file.readlines()]
    except FileNotFoundError:
        pytest.skip("whitelisted_bedrock_models.txt not found")

    _enforce_bedrock_converse_models(
        model_cost=litellm.model_cost, whitelist_models=whitelist_models
    )


@pytest.mark.flaky(retries=6, delay=2)
def test_model_info_bedrock_converse_enforcement(monkeypatch):
    """
    Test the enforcement of the whitelist by adding a fake model and ensuring the test fails.
    """
    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
    litellm.model_cost = litellm.get_model_cost_map(url="")

    # Add a fake unwhitelisted model
    litellm.model_cost["fake.bedrock-chat-model"] = {
        "litellm_provider": "bedrock",
        "mode": "chat",
    }

    try:
        # Load whitelist models from file
        with open("whitelisted_bedrock_models.txt", "r") as file:
            whitelist_models = [line.strip() for line in file.readlines()]

        # Check for unwhitelisted models
        with pytest.raises(AssertionError, match=r"fake\.bedrock-chat-model"):
            _enforce_bedrock_converse_models(
                model_cost=litellm.model_cost, whitelist_models=whitelist_models
            )
    except FileNotFoundError as e:
        pytest.skip("whitelisted_bedrock_models.txt not found")


@pytest.mark.parametrize("region", ("us-gov-east-1", "us-gov-west-1"))
@pytest.mark.parametrize("base_provider", ("bedrock_converse", "bedrock"))
def test_regional_bedrock_alias_requires_canonical_converse_metadata(
    region: str, base_provider: Literal["bedrock_converse", "bedrock"]
) -> None:
    base_model: Final = next(
        model for model in sorted(litellm.bedrock_converse_models) if BedrockModelInfo.get_base_model(model) == model
    )
    model: Final = f"bedrock/{region}/{base_model}"
    model_cost: Final[Mapping[str, ModelInfoBase]] = {
        model: {"litellm_provider": "bedrock", "mode": "chat"},
        base_model: {"litellm_provider": base_provider, "mode": "chat"},
    }
    assert BedrockModelInfo.get_bedrock_route(model) == "converse"
    if base_provider == "bedrock":
        with pytest.raises(AssertionError, match=re.escape(model)):
            _enforce_bedrock_converse_models(model_cost, ())
        return
    _enforce_bedrock_converse_models(model_cost, ())


GOVCLOUD_LAGGING_CAPABILITIES: Final = frozenset({"supports_bedrock_native_structured_output"})


def _govcloud_may_lag(model_key: str, capability: str) -> bool:
    is_govcloud: Final = model_key.startswith("us-gov.") or any(
        part.startswith("us-gov-") for part in model_key.split("/")
    )
    return is_govcloud and capability in GOVCLOUD_LAGGING_CAPABILITIES


def test_get_model_info_bedrock_models():
    """
    Check for drift in base model info for bedrock models and regional model info for bedrock models.
    GovCloud regions may lag their commercial base on Bedrock native structured output, never lead it.
    """
    from litellm.llms.bedrock.common_utils import BedrockModelInfo

    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    for k, v in litellm.model_cost.items():
        if v["litellm_provider"] == "bedrock":
            k = k.replace("*/", "")
            potential_commitments = [
                "1-month-commitment",
                "3-month-commitment",
                "6-month-commitment",
            ]
            if any(commitment in k for commitment in potential_commitments):
                for commitment in potential_commitments:
                    k = k.replace(f"{commitment}/", "")
            base_model = BedrockModelInfo.get_base_model(k)
            # get_base_model() returns model id without "bedrock/" prefix; cost map keys use "bedrock/<model>"
            base_model_key = (
                base_model
                if base_model in litellm.model_cost
                else f"bedrock/{base_model}"
            )
            if base_model_key not in litellm.model_cost:
                continue
            base_model_info = litellm.model_cost[base_model_key]
            for base_model_key, base_model_value in base_model_info.items():
                if "invoke/" in k:
                    continue
                if base_model_key.startswith("supports_"):
                    assert (
                        base_model_key in v
                    ), f"{base_model_key} is not in model cost map for {k}"
                    if _govcloud_may_lag(k, base_model_key):
                        assert base_model_value or not v[base_model_key], f"{base_model_key} is on {k} but not its base"
                        continue
                    assert (
                        v[base_model_key] == base_model_value
                    ), f"{base_model_key} is not equal to {base_model_value} for model {k}"


def _cross_region_base_model_key(
    model_key: str,
    prefixes: tuple[str, ...],
    model_cost: Mapping[str, Mapping[str, object]],
) -> str | None:
    matched_prefix: Final = next((prefix for prefix in prefixes if model_key.startswith(prefix)), None)
    if matched_prefix is None:
        return None
    base_model_key: Final = model_key[len(matched_prefix) :]
    return base_model_key if base_model_key in model_cost else None


def _cross_region_profiles(
    model_cost: Mapping[str, Mapping[str, object]],
    prefixes: tuple[str, ...],
) -> tuple[tuple[str, Mapping[str, object], str], ...]:
    return tuple(
        (model_key, model_info, base_model_key)
        for model_key, model_info in model_cost.items()
        if str(model_info.get("litellm_provider", "")).startswith("bedrock")
        and (
            base_model_key := _cross_region_base_model_key(model_key, prefixes, model_cost)
        )
        is not None
    )


def test_get_model_info_bedrock_cross_region_capability_parity():
    """
    Cross-region inference profiles carry litellm_provider "bedrock_converse", so the
    regional drift check above (which filters on "bedrock") never reaches them.
    GovCloud is a separate AWS partition that gets features later, so a us-gov profile
    may lack Bedrock native structured output its commercial base has, never the reverse.
    """
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    prefixes = ("us.", "eu.", "apac.", "us-gov.")
    checked = 0

    for k, v in litellm.model_cost.items():
        if not str(v.get("litellm_provider", "")).startswith("bedrock"):
            continue
        base_model_key = next(
            (k[len(p) :] for p in prefixes if k.startswith(p)),
            None,
        )
        if base_model_key is None or base_model_key not in litellm.model_cost:
            continue
        checked += 1
        for cap, base_value in litellm.model_cost[base_model_key].items():
            if not cap.startswith("supports_"):
                continue
            assert cap in v, f"{cap} is on {base_model_key} but missing from {k}"
            if _govcloud_may_lag(k, cap):
                assert base_value or not v[cap], f"{cap} is on {k} but not on {base_model_key}"
                continue
            assert (
                v[cap] == base_value
            ), f"{cap} is {v[cap]} on {k} but {base_value} on {base_model_key}"

    assert checked > 0, "no cross-region bedrock profiles found - the filter is inert"


def _is_positive_cost(value: object) -> bool:
    return isinstance(value, (int, float)) and value > 0


def test_get_model_info_bedrock_priced_cross_region_profile_has_priced_base():
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    prefixes = ("us.", "eu.", "apac.", "us-gov.", "au.", "global.")
    checked = 0

    for k, v in litellm.model_cost.items():
        if not str(v.get("litellm_provider", "")).startswith("bedrock"):
            continue
        base_model_key = next(
            (k[len(p) :] for p in prefixes if k.startswith(p)),
            None,
        )
        if base_model_key is None or base_model_key not in litellm.model_cost:
            continue
        checked += 1
        base = litellm.model_cost[base_model_key]
        for cost_key in ("input_cost_per_token", "output_cost_per_token"):
            if (v.get(cost_key) or 0) > 0:
                assert (
                    base.get(cost_key) or 0
                ) > 0, f"{k} charges {cost_key} but its base {base_model_key} is free"

    assert checked > 0, "no cross-region bedrock profiles found - the filter is inert"


def test_get_model_info_case_insensitive_lookup(monkeypatch):
    """
    Test that model info lookup is case-insensitive.

    This ensures that users can use lowercase model names even when the model cost
    map has mixed-case keys (e.g., "Qwen/Qwen3-Next-80B-A3B-Thinking").

    Related Slack discussion: Users were getting "does not support parameters: ['tools']"
    errors when using lowercase model names like "qwen/qwen3-next-80b-a3b-thinking"
    because the lookup was case-sensitive.
    """
    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
    litellm.model_cost = litellm.get_model_cost_map(url="")

    # Register a test model with mixed-case name
    litellm.register_model(
        {
            "together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": {
                "input_cost_per_token": 0.0001,
                "output_cost_per_token": 0.0002,
                "litellm_provider": "together_ai",
                "supports_function_calling": True,
            }
        }
    )

    # Test 1: Exact case should work
    info = litellm.get_model_info(
        model="Qwen/Qwen3-Next-80B-A3B-Thinking", custom_llm_provider="together_ai"
    )
    assert info is not None
    assert info["supports_function_calling"] is True

    # Test 2: Lowercase should also work (case-insensitive lookup)
    info_lower = litellm.get_model_info(
        model="qwen/qwen3-next-80b-a3b-thinking", custom_llm_provider="together_ai"
    )
    assert info_lower is not None
    assert info_lower["supports_function_calling"] is True

    # Test 3: Mixed case should also work
    info_mixed = litellm.get_model_info(
        model="QWEN/qwen3-NEXT-80b-a3b-thinking", custom_llm_provider="together_ai"
    )
    assert info_mixed is not None
    assert info_mixed["supports_function_calling"] is True


def test_get_model_info_case_insensitive_supports_function_calling(monkeypatch):
    """
    Test that supports_function_calling check works with case-insensitive model lookup.
    """
    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
    litellm.model_cost = litellm.get_model_cost_map(url="")

    # Register a model with mixed-case name that supports function calling
    litellm.register_model(
        {
            "test_provider/TestModel-ABC": {
                "input_cost_per_token": 0.0001,
                "output_cost_per_token": 0.0002,
                "litellm_provider": "test_provider",
                "supports_function_calling": True,
            }
        }
    )

    # Test that supports_function_calling works with lowercase model name
    from litellm.utils import supports_function_calling

    # Exact case
    assert (
        supports_function_calling("TestModel-ABC", custom_llm_provider="test_provider")
        is True
    )

    # Lowercase (should now work with case-insensitive lookup)
    assert (
        supports_function_calling("testmodel-abc", custom_llm_provider="test_provider")
        is True
    )


def test_get_model_info_custom_model_router():
    from litellm import Router
    from litellm import get_model_info

    litellm.turn_on_debug()

    router = Router(
        model_list=[
            {
                "model_name": "ma-summary",
                "litellm_params": {
                    "api_base": "http://ma-mix-llm-serving.cicero.svc.cluster.local/v1",
                    "input_cost_per_token": 1,
                    "output_cost_per_token": 1,
                    "model": "openai/meta-llama/Meta-Llama-3-8B-Instruct",
                },
                "model_info": {
                    "id": "c20d603e-1166-4e0f-aa65-ed9c476ad4ca",
                },
            }
        ]
    )
    info = get_model_info("c20d603e-1166-4e0f-aa65-ed9c476ad4ca")
    print("info", info)
    assert info is not None


def test_get_model_info_custom_provider():
    # Custom provider example copied from https://docs.litellm.ai/docs/providers/custom_llm_server:
    import litellm
    from litellm import CustomLLM, completion

    class MyCustomLLM(CustomLLM):
        def completion(self, *args, **kwargs) -> litellm.ModelResponse:
            return litellm.completion(
                model="gpt-3.5-turbo",
                messages=[{"role": "user", "content": "Hello world"}],
                mock_response="Hi!",
            )  # type: ignore

    my_custom_llm = MyCustomLLM()

    litellm.custom_provider_map = [  # 👈 KEY STEP - REGISTER HANDLER
        {"provider": "my-custom-llm", "custom_handler": my_custom_llm}
    ]

    resp = completion(
        model="my-custom-llm/my-fake-model",
        messages=[{"role": "user", "content": "Hello world!"}],
    )

    assert resp.choices[0].message.content == "Hi!"

    # Register model info
    model_info = {"my-custom-llm/my-fake-model": {"max_tokens": 2048}}
    litellm.register_model(model_info)

    # Get registered model info
    from litellm import get_model_info

    get_model_info(
        model="my-custom-llm/my-fake-model"
    )  # 💥 "Exception: This model isn't mapped yet." in v1.56.10


def test_get_model_info_huggingface_models(monkeypatch):
    from litellm import Router
    from litellm.types.router import ModelGroupInfo

    monkeypatch.setenv("HUGGINGFACE_API_KEY", "hf_abc123")

    with respx.mock(assert_all_called=False) as huggingface:
        huggingface.get(host="huggingface.co").mock(side_effect=httpx.ConnectError("huggingface.co is unreachable"))
        router = Router(
            model_list=[
                {
                    "model_name": "meta-llama/Meta-Llama-3-8B-Instruct",
                    "litellm_params": {
                        "model": "huggingface/meta-llama/Meta-Llama-3-8B-Instruct",
                        "api_base": "https://router.huggingface.co/hf-inference/models/meta-llama/Meta-Llama-3-8B-Instruct",
                        "api_key": os.environ["HUGGINGFACE_API_KEY"],
                    },
                }
            ]
        )
        info = litellm.get_model_info("huggingface/meta-llama/Meta-Llama-3-8B-Instruct")
    assert info is not None

    ModelGroupInfo(
        model_group="meta-llama/Meta-Llama-3-8B-Instruct",
        providers=["huggingface"],
        **info,
    )