diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c4f2b45f38..d168290e8a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/tests/test_litellm/proxy/test_empty_model_list.py b/tests/test_litellm/proxy/test_empty_model_list.py index 6b3e59d319..dd900d3eb5 100644 --- a/tests/test_litellm/proxy/test_empty_model_list.py +++ b/tests/test_litellm/proxy/test_empty_model_list.py @@ -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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d8970a76a9..f1854380ef 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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