This commit is contained in:
Utkash Dubey 2025-03-03 04:16:12 -08:00
parent db83cbe5c0
commit fa88bc9632
19 changed files with 191 additions and 141 deletions

View file

@ -58,16 +58,14 @@ def test_get_model_info_shows_correct_supports_vision():
def test_get_model_info_shows_assistant_prefill():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm.model_cost = litellm.get_locally_cached_model_cost_map()
info = litellm.get_model_info("deepseek/deepseek-chat")
print("info", info)
assert info.get("supports_assistant_prefill") is True
def test_get_model_info_shows_supports_prompt_caching():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm.model_cost = litellm.get_locally_cached_model_cost_map()
info = litellm.get_model_info("deepseek/deepseek-chat")
print("info", info)
assert info.get("supports_prompt_caching") is True
@ -116,8 +114,7 @@ def test_get_model_info_gemini():
"""
Tests if ALL gemini models have 'tpm' and 'rpm' in the model info
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm.model_cost = litellm.get_locally_cached_model_cost_map()
model_map = litellm.model_cost
for model, info in model_map.items():
@ -127,8 +124,7 @@ def test_get_model_info_gemini():
def test_get_model_info_bedrock_region():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm.model_cost = litellm.get_locally_cached_model_cost_map()
args = {
"model": "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
"custom_llm_provider": "bedrock",
@ -212,8 +208,7 @@ def test_model_info_bedrock_converse(monkeypatch):
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="")
litellm.model_cost = litellm.get_locally_cached_model_cost_map()
try:
# Load whitelist models from file
with open("whitelisted_bedrock_models.txt", "r") as file:
@ -231,8 +226,7 @@ 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="")
litellm.model_cost = litellm.get_locally_cached_model_cost_map()
# Add a fake unwhitelisted model
litellm.model_cost["fake.bedrock-chat-model"] = {
@ -323,8 +317,7 @@ def test_get_model_info_bedrock_models():
"""
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="")
litellm.model_cost = litellm.get_locally_cached_model_cost_map()
for k, v in litellm.model_cost.items():
if v["litellm_provider"] == "bedrock":