mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(bedrock): preserve deployment isolation on older release interfaces
This commit is contained in:
Joshua Valluru
parent
df1541a0a2
commit
97c0f49cac
5 changed files with 39 additions and 504 deletions
|
|
@ -2695,3 +2695,15 @@ def text_completion_prompt_to_messages(prompt: object) -> tuple[AllMessageValues
|
|||
if string_entries and len(string_entries) == len(entries):
|
||||
return tuple(ChatCompletionUserMessage(role="user", content=entry) for entry in string_entries)
|
||||
raise ValueError(f"`prompt` must be a non-empty string or a non-empty list of strings. Got: {prompt_type_name}.")
|
||||
|
||||
|
||||
_REGEX_LOOKAROUND_RE: Final = re.compile(r"\(\?<?[=!]")
|
||||
|
||||
|
||||
def _uses_regex_lookaround(pattern: str) -> bool:
|
||||
return _REGEX_LOOKAROUND_RE.search(pattern) is not None
|
||||
|
||||
|
||||
def _schema_at(container: Mapping[str, object], key: str) -> Mapping[str, object] | None:
|
||||
value: Final = container.get(key)
|
||||
return value if isinstance(value, dict) else None
|
||||
|
|
|
|||
|
|
@ -3872,6 +3872,7 @@ SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = (
|
|||
frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__)
|
||||
- frozenset(CustomPricingLiteLLMParams.model_fields)
|
||||
- DEPLOYMENT_SCOPED_PRICING_FIELDS
|
||||
- DEPLOYMENT_SCOPED_CAPABILITY_FIELDS
|
||||
- DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from typing import Final
|
||||
# What is this?
|
||||
## Unit testing for the 'get_model_info()' function
|
||||
import os
|
||||
|
|
@ -279,6 +280,7 @@ def test_get_model_info_custom_model_router():
|
|||
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
|
||||
|
||||
|
|
@ -298,11 +300,7 @@ def test_get_model_info_bedrock_models():
|
|||
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}"
|
||||
)
|
||||
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]
|
||||
|
|
@ -310,18 +308,21 @@ def test_get_model_info_bedrock_models():
|
|||
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}"
|
||||
assert (
|
||||
v[base_model_key] == base_model_value
|
||||
), f"{base_model_key} is not equal to {base_model_value} for model {k}"
|
||||
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 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="")
|
||||
|
|
@ -343,9 +344,10 @@ def test_get_model_info_bedrock_cross_region_capability_parity():
|
|||
if not cap.startswith("supports_"):
|
||||
continue
|
||||
assert cap in v, f"{cap} is on {base_model_key} but missing from {k}"
|
||||
assert (
|
||||
v[cap] == base_value
|
||||
), f"{cap} is {v[cap]} on {k} but {base_value} on {base_model_key}"
|
||||
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"
|
||||
|
||||
|
|
@ -487,3 +489,12 @@ def test_get_model_info_case_insensitive_supports_function_calling(monkeypatch):
|
|||
supports_function_calling("testmodel-abc", custom_llm_provider="test_provider")
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import litellm
|
||||
"""Tests for `BedrockConverseLLM.completion`.
|
||||
|
||||
AWS credential resolution is stubbed so nothing reaches STS.
|
||||
|
|
|
|||
|
|
@ -1,490 +0,0 @@
|
|||
import json
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Collection, Iterator, Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal, cast
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
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
|
||||
|
||||
|
||||
@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, 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
|
||||
Loading…
Add table
Reference in a new issue