tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-10 22:24:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixing tests and linting

Browse Source
This commit is contained in:
yuneng-jiang
2026-01-21 11:02:39 -08:00
parent b5a7d2ab34
commit d0e35751a1
3 changed files with 312 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/proxy_server.py
+74 -52
View File
@@ -7503,6 +7503,77 @@ async def get_all_team_and_direct_access_models(
return all_models
def _enrich_model_info_with_litellm_data(
model: Dict[str, Any], debug: bool = False, llm_router: Optional[Router] = None
) -> Dict[str, Any]:
"""
Enrich a model dictionary with litellm model info (pricing, context window, etc.)
and remove sensitive information.
Args:
model: Model dictionary to enrich
debug: Whether to include debug information like openai_client
llm_router: Optional router instance for debug info
Returns:
Enriched model dictionary with sensitive info removed
"""
# provided model_info in config.yaml
model_info = model.get("model_info", {})
if debug is True:
_openai_client = "None"
if llm_router is not None:
_openai_client = (
llm_router._get_client(
deployment=model, kwargs={}, client_type="async"
)
or "None"
)
else:
_openai_client = "llm_router_is_None"
openai_client = str(_openai_client)
model["openai_client"] = openai_client
# read litellm model_prices_and_context_window.json to get the following:
# input_cost_per_token, output_cost_per_token, max_tokens
litellm_model_info = get_litellm_model_info(model=model)
# 2nd pass on the model, try seeing if we can find model in litellm model_cost map
if litellm_model_info == {}:
# use litellm_param model_name to get model_info
litellm_params = model.get("litellm_params", {})
litellm_model = litellm_params.get("model", None)
try:
litellm_model_info = litellm.get_model_info(model=litellm_model)
except Exception:
litellm_model_info = {}
# 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map
if litellm_model_info == {}:
# use litellm_param model_name to get model_info
litellm_params = model.get("litellm_params", {})
litellm_model = litellm_params.get("model", None)
if litellm_model:
split_model = litellm_model.split("/")
if len(split_model) > 0:
litellm_model = split_model[-1]
try:
litellm_model_info = litellm.get_model_info(
model=litellm_model, custom_llm_provider=split_model[0]
)
except Exception:
litellm_model_info = {}
for k, v in litellm_model_info.items():
if k not in model_info:
model_info[k] = v
model["model_info"] = model_info
# don't return the api key / vertex credentials
# don't return the llm credentials
model = remove_sensitive_info_from_deployment(
model, excluded_keys={"litellm_credential_name"}
)
return model
@router.get(
"/v2/model/info",
description="v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true",
@@ -7573,58 +7644,9 @@ async def model_info_v2(
all_models=all_models,
)
# fill in model info based on config.yaml and litellm model_prices_and_context_window.json
for _model in all_models:
# provided model_info in config.yaml
model_info = _model.get("model_info", {})
if debug is True:
_openai_client = "None"
if llm_router is not None:
_openai_client = (
llm_router._get_client(
deployment=_model, kwargs={}, client_type="async"
)
or "None"
)
else:
_openai_client = "llm_router_is_None"
openai_client = str(_openai_client)
_model["openai_client"] = openai_client
# read litellm model_prices_and_context_window.json to get the following:
# input_cost_per_token, output_cost_per_token, max_tokens
litellm_model_info = get_litellm_model_info(model=_model)
# 2nd pass on the model, try seeing if we can find model in litellm model_cost map
if litellm_model_info == {}:
# use litellm_param model_name to get model_info
litellm_params = _model.get("litellm_params", {})
litellm_model = litellm_params.get("model", None)
try:
litellm_model_info = litellm.get_model_info(model=litellm_model)
except Exception:
litellm_model_info = {}
# 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map
if litellm_model_info == {}:
# use litellm_param model_name to get model_info
litellm_params = _model.get("litellm_params", {})
litellm_model = litellm_params.get("model", None)
split_model = litellm_model.split("/")
if len(split_model) > 0:
litellm_model = split_model[-1]
try:
litellm_model_info = litellm.get_model_info(
model=litellm_model, custom_llm_provider=split_model[0]
)
except Exception:
litellm_model_info = {}
for k, v in litellm_model_info.items():
if k not in model_info:
model_info[k] = v
_model["model_info"] = model_info
# don't return the api key / vertex credentials
# don't return the llm credentials
_model = remove_sensitive_info_from_deployment(
_model, excluded_keys={"litellm_credential_name"}
for i, _model in enumerate(all_models):
all_models[i] = _enrich_model_info_with_litellm_data(
model=_model, debug=debug, llm_router=llm_router
)
verbose_proxy_logger.debug("all_models: %s", all_models)
tests/test_litellm/proxy/test_empty_model_list.py
+54 -4
View File
@@ -32,7 +32,7 @@ class TestEmptyModelListHandling:
self, client, monkeypatch
):
"""
Test that /v2/model/info returns {"data": []} instead of 500
Test that /v2/model/info returns paginated empty response instead of 500
when llm_router is None.
"""
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
@@ -56,13 +56,18 @@ class TestEmptyModelListHandling:
)
assert response.status_code == 200
assert response.json() == {"data": []}
data = response.json()
assert data["data"] == []
assert data["total_count"] == 0
assert data["current_page"] == 1
assert data["total_pages"] == 0
assert data["size"] == 50 # default page size
def test_v2_model_info_returns_empty_data_when_model_list_empty(
self, client, monkeypatch
):
"""
Test that /v2/model/info returns {"data": []} instead of 500
Test that /v2/model/info returns paginated empty response instead of 500
when llm_router exists but model_list is empty.
"""
mock_router = MagicMock()
@@ -89,7 +94,52 @@ class TestEmptyModelListHandling:
)
assert response.status_code == 200
assert response.json() == {"data": []}
data = response.json()
assert data["data"] == []
assert data["total_count"] == 0
assert data["current_page"] == 1
assert data["total_pages"] == 0
assert data["size"] == 50 # default page size
def test_v2_model_info_pagination_with_empty_results(
self, client, monkeypatch
):
"""
Test that /v2/model/info pagination parameters work correctly
when there are no models (empty results).
"""
mock_router = MagicMock()
mock_router.model_list = []
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", [])
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
with patch(
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth",
return_value=MagicMock(
user_id="test-user",
team_id=None,
team_models=[],
models=[],
user_role="proxy_admin",
),
):
# Test with custom pagination parameters
response = client.get(
"/v2/model/info",
params={"page": 2, "size": 25},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
data = response.json()
assert data["data"] == []
assert data["total_count"] == 0
assert data["current_page"] == 2 # Should respect the page parameter
assert data["total_pages"] == 0
assert data["size"] == 25 # Should respect the size parameter
def test_model_group_info_returns_empty_data_when_model_list_none(
self, client, monkeypatch
tests/test_litellm/proxy/test_proxy_server.py
+184
View File
@@ -3441,3 +3441,187 @@ async def test_model_info_v2_pagination_edge_cases(monkeypatch):
finally:
app.dependency_overrides = original_overrides
def test_enrich_model_info_with_litellm_data():
"""
Test the _enrich_model_info_with_litellm_data helper function.
Tests model info enrichment, debug mode, and sensitive info removal.
"""
from unittest.mock import MagicMock, patch
from litellm.proxy.proxy_server import _enrich_model_info_with_litellm_data
# Test Case 1: Basic model enrichment without debug
model = {
"model_name": "test-model",
"litellm_params": {"model": "gpt-3.5-turbo"},
"model_info": {"id": "test-model"},
"api_key": "sk-secret-key", # Should be removed
}
with patch("litellm.proxy.proxy_server.get_litellm_model_info") as mock_get_info, patch(
"litellm.proxy.proxy_server.remove_sensitive_info_from_deployment"
) as mock_remove_sensitive:
mock_get_info.return_value = {
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
"max_tokens": 4096,
}
mock_remove_sensitive.return_value = {
"model_name": "test-model",
"litellm_params": {"model": "gpt-3.5-turbo"},
"model_info": {
"id": "test-model",
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
"max_tokens": 4096,
},
}
result = _enrich_model_info_with_litellm_data(model=model, debug=False)
# Verify get_litellm_model_info was called
mock_get_info.assert_called_once_with(model=model)
# Verify remove_sensitive_info_from_deployment was called
mock_remove_sensitive.assert_called_once()
# Verify result doesn't have api_key
assert "api_key" not in result
# Verify model_info was enriched
assert "input_cost_per_token" in result["model_info"]
# Test Case 2: Model enrichment with debug mode
model_with_debug = {
"model_name": "test-model-debug",
"litellm_params": {"model": "gpt-4"},
"model_info": {},
}
mock_router = MagicMock()
mock_client = MagicMock()
mock_router._get_client.return_value = mock_client
with patch("litellm.proxy.proxy_server.get_litellm_model_info") as mock_get_info, patch(
"litellm.proxy.proxy_server.remove_sensitive_info_from_deployment"
) as mock_remove_sensitive:
mock_get_info.return_value = {}
mock_remove_sensitive.return_value = {
"model_name": "test-model-debug",
"litellm_params": {"model": "gpt-4"},
"model_info": {},
"openai_client": str(mock_client),
}
result = _enrich_model_info_with_litellm_data(
model=model_with_debug, debug=True, llm_router=mock_router
)
# Verify debug info was added
mock_remove_sensitive.assert_called_once()
call_args = mock_remove_sensitive.call_args[0][0]
assert "openai_client" in call_args
# Verify router._get_client was called for debug
mock_router._get_client.assert_called_once()
# Test Case 3: Model with fallback to litellm.get_model_info
model_fallback = {
"model_name": "test-model-fallback",
"litellm_params": {"model": "claude-3-opus"},
"model_info": {},
}
with patch("litellm.proxy.proxy_server.get_litellm_model_info") as mock_get_info, patch(
"litellm.get_model_info"
) as mock_litellm_info, patch(
"litellm.proxy.proxy_server.remove_sensitive_info_from_deployment"
) as mock_remove_sensitive:
# First call returns empty, triggering fallback
mock_get_info.return_value = {}
mock_litellm_info.return_value = {
"input_cost_per_token": 0.015,
"output_cost_per_token": 0.075,
"max_tokens": 200000,
}
mock_remove_sensitive.return_value = {
"model_name": "test-model-fallback",
"litellm_params": {"model": "claude-3-opus"},
"model_info": {
"input_cost_per_token": 0.015,
"output_cost_per_token": 0.075,
"max_tokens": 200000,
},
}
result = _enrich_model_info_with_litellm_data(model=model_fallback, debug=False)
# Verify fallback was attempted
mock_litellm_info.assert_called_once_with(model="claude-3-opus")
# Verify model_info was enriched with fallback data
call_args = mock_remove_sensitive.call_args[0][0]
assert call_args["model_info"]["input_cost_per_token"] == 0.015
# Test Case 4: Model with split model name fallback
model_split = {
"model_name": "test-model-split",
"litellm_params": {"model": "azure/gpt-4"},
"model_info": {},
}
with patch("litellm.proxy.proxy_server.get_litellm_model_info") as mock_get_info, patch(
"litellm.get_model_info"
) as mock_litellm_info, patch(
"litellm.proxy.proxy_server.remove_sensitive_info_from_deployment"
) as mock_remove_sensitive:
# Both first and second pass return empty, triggering third pass
mock_get_info.return_value = {}
# Second pass (no split)
mock_litellm_info.side_effect = [
{}, # First call returns empty
{"max_tokens": 8192}, # Third pass with split succeeds
]
mock_remove_sensitive.return_value = {
"model_name": "test-model-split",
"litellm_params": {"model": "azure/gpt-4"},
"model_info": {"max_tokens": 8192},
}
result = _enrich_model_info_with_litellm_data(model=model_split, debug=False)
# Verify third pass was attempted with split model name
assert mock_litellm_info.call_count == 2
# Check that second call used split model name
second_call = mock_litellm_info.call_args_list[1]
assert second_call[1]["model"] == "gpt-4"
assert second_call[1]["custom_llm_provider"] == "azure"
# Test Case 5: Model with existing model_info (should preserve existing keys)
model_existing = {
"model_name": "test-model-existing",
"litellm_params": {"model": "gpt-3.5-turbo"},
"model_info": {"id": "existing-id", "custom_key": "custom_value"},
}
with patch("litellm.proxy.proxy_server.get_litellm_model_info") as mock_get_info, patch(
"litellm.proxy.proxy_server.remove_sensitive_info_from_deployment"
) as mock_remove_sensitive:
mock_get_info.return_value = {
"input_cost_per_token": 0.001,
"id": "new-id", # Should not override existing "id"
}
mock_remove_sensitive.return_value = {
"model_name": "test-model-existing",
"litellm_params": {"model": "gpt-3.5-turbo"},
"model_info": {
"id": "existing-id", # Existing key preserved
"custom_key": "custom_value", # Existing key preserved
"input_cost_per_token": 0.001, # New key added
},
}
result = _enrich_model_info_with_litellm_data(model=model_existing, debug=False)
# Verify existing keys are preserved
call_args = mock_remove_sensitive.call_args[0][0]
assert call_args["model_info"]["id"] == "existing-id"
assert call_args["model_info"]["custom_key"] == "custom_value"
assert call_args["model_info"]["input_cost_per_token"] == 0.001