From 762a3ef090df2fcce7130a6454403c04598d1f80 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 13 Jan 2026 16:39:57 +0530 Subject: [PATCH] Add support for 0 cost models --- litellm/proxy/auth/auth_checks.py | 179 ++++-- litellm/proxy/auth/user_api_key_auth.py | 100 ++- .../test_zero_cost_model_budget_bypass.py | 590 ++++++++++++++++++ 3 files changed, 782 insertions(+), 87 deletions(-) create mode 100644 tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a741869e5f..1879b30625 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -74,6 +74,75 @@ db_cache_expiry = DEFAULT_IN_MEMORY_TTL # refresh every 5s all_routes = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value +def _is_model_cost_zero( + model: Optional[Union[str, List[str]]], llm_router: Optional[Router] +) -> bool: + """ + Check if a model has zero cost (no configured pricing). + + Uses the router's get_model_group_info method to get pricing information. + + Args: + model: The model name or list of model names + llm_router: The LiteLLM router instance + + Returns: + bool: True if all costs for the model are zero, False otherwise + """ + if model is None or llm_router is None: + return False + + # Handle list of models + model_list = [model] if isinstance(model, str) else model + + for model_name in model_list: + try: + # Use router's get_model_group_info method directly for better reliability + model_group_info = llm_router.get_model_group_info(model_group=model_name) + + if model_group_info is None: + # Model not found or no pricing info available + # Conservative approach: assume it has cost + verbose_proxy_logger.debug( + f"No model group info found for {model_name}, assuming it has cost" + ) + return False + + # Check costs for this model + # Only allow bypass if BOTH costs are explicitly set to 0 (not None) + input_cost = model_group_info.input_cost_per_token + output_cost = model_group_info.output_cost_per_token + + # If costs are not explicitly configured (None), assume it has cost + if input_cost is None or output_cost is None: + verbose_proxy_logger.debug( + f"Model {model_name} has undefined cost (input: {input_cost}, output: {output_cost}), assuming it has cost" + ) + return False + + # If either cost is non-zero, return False + if input_cost > 0 or output_cost > 0: + verbose_proxy_logger.debug( + f"Model {model_name} has non-zero cost (input: {input_cost}, output: {output_cost})" + ) + return False + + # This model has zero cost explicitly configured + verbose_proxy_logger.debug( + f"Model {model_name} has zero cost explicitly configured (input: {input_cost}, output: {output_cost})" + ) + + except Exception as e: + # If we can't determine the cost, assume it has cost (conservative approach) + verbose_proxy_logger.debug( + f"Error checking cost for model {model_name}: {str(e)}, assuming it has cost" + ) + return False + + # All models checked have zero cost + return True + + async def common_checks( request_body: dict, team_object: Optional[LiteLLM_TeamTable], @@ -86,6 +155,7 @@ async def common_checks( proxy_logging_obj: ProxyLogging, valid_token: Optional[UserAPIKeyAuth], request: Request, + skip_budget_checks: bool = False, ) -> bool: """ Common checks across jwt + key-based auth. @@ -137,64 +207,66 @@ async def common_checks( user_object=user_object, ) - # 3. If team is in budget - await _team_max_budget_check( - team_object=team_object, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ) + # If this is a free model, skip all budget checks + if not skip_budget_checks: + # 3. If team is in budget + await _team_max_budget_check( + team_object=team_object, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + ) - # 3.1. If organization is in budget - await _organization_max_budget_check( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + # 3.1. If organization is in budget + await _organization_max_budget_check( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) - await _tag_max_budget_check( - request_body=request_body, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ) + await _tag_max_budget_check( + request_body=request_body, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + ) - # 4. If user is in budget - ## 4.1 check personal budget, if personal key - if ( - (team_object is None or team_object.team_id is None) - and user_object is not None - and user_object.max_budget is not None - ): - user_budget = user_object.max_budget - if user_budget < user_object.spend: - raise litellm.BudgetExceededError( - current_cost=user_object.spend, - max_budget=user_budget, - message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}", - ) + # 4. If user is in budget + ## 4.1 check personal budget, if personal key + if ( + (team_object is None or team_object.team_id is None) + and user_object is not None + and user_object.max_budget is not None + ): + user_budget = user_object.max_budget + if user_budget < user_object.spend: + raise litellm.BudgetExceededError( + current_cost=user_object.spend, + max_budget=user_budget, + message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}", + ) - ## 4.2 check team member budget, if team key - await _check_team_member_budget( - team_object=team_object, - user_object=user_object, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + ## 4.2 check team member budget, if team key + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) - # 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget - if end_user_object is not None and end_user_object.litellm_budget_table is not None: - end_user_budget = end_user_object.litellm_budget_table.max_budget - if end_user_budget is not None and end_user_object.spend > end_user_budget: - raise litellm.BudgetExceededError( - current_cost=end_user_object.spend, - max_budget=end_user_budget, - message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}", - ) + # 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget + if end_user_object is not None and end_user_object.litellm_budget_table is not None: + end_user_budget = end_user_object.litellm_budget_table.max_budget + if end_user_budget is not None and end_user_object.spend > end_user_budget: + raise litellm.BudgetExceededError( + current_cost=end_user_object.spend, + max_budget=end_user_budget, + message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}", + ) # 6. [OPTIONAL] If 'enforce_user_param' enabled - did developer pass in 'user' param for openai endpoints if ( @@ -237,6 +309,7 @@ async def common_checks( # 7. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget if ( litellm.max_budget > 0 + and not skip_budget_checks and global_proxy_spend is not None # only run global budget checks for OpenAI routes # Reason - the Admin UI should continue working if the proxy crosses it's global budget diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 401aa7fd44..efac74219d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -586,6 +586,21 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_object is not None else None, ) + + # Check if model has zero cost - if so, skip all budget checks + model = get_model_from_request(request_data, route) + skip_budget_checks = False + if model is not None and llm_router is not None: + from litellm.proxy.auth.auth_checks import _is_model_cost_zero + + skip_budget_checks = _is_model_cost_zero( + model=model, llm_router=llm_router + ) + if skip_budget_checks: + verbose_proxy_logger.info( + f"Skipping all budget checks for zero-cost model: {model}" + ) + # run through common checks _ = await common_checks( request=request, @@ -599,6 +614,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 llm_router=llm_router, proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, + skip_budget_checks=skip_budget_checks, ) # return UserAPIKeyAuth object @@ -990,8 +1006,22 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) user_obj = None + # Check 2a. Check if model has zero cost - if so, skip all budget checks + model = get_model_from_request(request_data, route) + skip_budget_checks = False + if model is not None and llm_router is not None: + from litellm.proxy.auth.auth_checks import _is_model_cost_zero + + skip_budget_checks = _is_model_cost_zero( + model=model, llm_router=llm_router + ) + if skip_budget_checks: + verbose_proxy_logger.info( + f"Skipping all budget checks for zero-cost model: {model}" + ) + # Check 3. Check if user is in their team budget - if valid_token.team_member_spend is not None: + if not skip_budget_checks and valid_token.team_member_spend is not None: if prisma_client is not None: _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" @@ -1055,46 +1085,47 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 param=abbreviate_api_key(api_key=api_key), ) - # Check 4. Token Spend is under budget - if RouteChecks.is_llm_api_route(route=route): - await _virtual_key_max_budget_check( + if not skip_budget_checks: + # Check 4. Token Spend is under budget + if RouteChecks.is_llm_api_route(route=route): + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=user_obj, + ) + + # Check 5. Max Budget Alert Check + await _virtual_key_max_budget_alert_check( valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, user_obj=user_obj, ) - # Check 5. Max Budget Alert Check - await _virtual_key_max_budget_alert_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # Check 6. Soft Budget Check - await _virtual_key_soft_budget_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # Check 5. Token Model Spend is under Model budget - max_budget_per_model = valid_token.model_max_budget - current_model = request_data.get("model", None) - - if ( - max_budget_per_model is not None - and isinstance(max_budget_per_model, dict) - and len(max_budget_per_model) > 0 - and prisma_client is not None - and current_model is not None - and valid_token.token is not None - ): - ## GET THE SPEND FOR THIS MODEL - await model_max_budget_limiter.is_key_within_model_budget( - user_api_key_dict=valid_token, - model=current_model, + # Check 6. Soft Budget Check + await _virtual_key_soft_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=user_obj, ) + # Check 5. Token Model Spend is under Model budget + max_budget_per_model = valid_token.model_max_budget + current_model = request_data.get("model", None) + + if ( + max_budget_per_model is not None + and isinstance(max_budget_per_model, dict) + and len(max_budget_per_model) > 0 + and prisma_client is not None + and current_model is not None + and valid_token.token is not None + ): + ## GET THE SPEND FOR THIS MODEL + await model_max_budget_limiter.is_key_within_model_budget( + user_api_key_dict=valid_token, + model=current_model, + ) + # Check 6: Additional Common Checks across jwt + key auth if valid_token.team_id is not None: _team_obj: Optional[LiteLLM_TeamTable] = LiteLLM_TeamTable( @@ -1162,6 +1193,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 llm_router=llm_router, proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, + skip_budget_checks=skip_budget_checks, ) # Token passed all checks if valid_token is None: diff --git a/tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py b/tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py new file mode 100644 index 0000000000..bc818fc0dc --- /dev/null +++ b/tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py @@ -0,0 +1,590 @@ +""" +Tests for zero-cost model budget bypass functionality. + +When a user exceeds their budget, the system should still allow requests +to models with zero cost (e.g., on-premises models). +""" + +import asyncio +from typing import Optional +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +from litellm.caching.caching import DualCache +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + _check_team_member_budget, + _is_model_cost_zero, + _team_max_budget_check, + common_checks, +) +from litellm.proxy.utils import ProxyLogging +from litellm.router import Router +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + +@pytest.fixture +def mock_router_with_zero_cost_model(): + """Create a mock router with a zero-cost model.""" + router = Router( + model_list=[ + { + "model_name": "on-prem-model", + "litellm_params": { + "model": "ollama/llama2", + "api_base": "http://localhost:11434", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": { + "id": "on-prem-model-id", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + { + "model_name": "cloud-model", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "sk-test", + }, + "model_info": { + "id": "cloud-model-id", + }, + }, + ] + ) + return router + + +@pytest.fixture +def mock_router_with_paid_model(): + """Create a mock router with only paid models.""" + router = Router( + model_list=[ + { + "model_name": "cloud-model", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "sk-test", + }, + "model_info": { + "id": "cloud-model-id", + }, + } + ] + ) + return router + + +@pytest.fixture +def mock_proxy_logging(): + """Create a mock ProxyLogging instance.""" + proxy_logging = ProxyLogging(user_api_key_cache=None) + + async def mock_budget_alerts(*args, **kwargs): + pass + + proxy_logging.budget_alerts = mock_budget_alerts + return proxy_logging + + +class TestIsModelCostZero: + """Tests for _is_model_cost_zero helper function.""" + + def test_zero_cost_model_in_router(self, mock_router_with_zero_cost_model): + """Test that a zero-cost model in router is correctly identified.""" + result = _is_model_cost_zero( + model="on-prem-model", llm_router=mock_router_with_zero_cost_model + ) + assert result is True + + def test_paid_model_in_router(self, mock_router_with_zero_cost_model): + """Test that a paid model is correctly identified as non-zero cost.""" + with patch("litellm.get_model_info") as mock_get_model_info: + # Mock the return value for gpt-3.5-turbo + mock_get_model_info.return_value = { + "input_cost_per_token": 0.0000015, + "output_cost_per_token": 0.000002, + } + result = _is_model_cost_zero( + model="cloud-model", llm_router=mock_router_with_zero_cost_model + ) + assert result is False + + def test_none_model(self, mock_router_with_zero_cost_model): + """Test that None model returns False.""" + result = _is_model_cost_zero( + model=None, llm_router=mock_router_with_zero_cost_model + ) + assert result is False + + def test_none_router(self): + """Test that None router returns False.""" + result = _is_model_cost_zero(model="some-model", llm_router=None) + assert result is False + + def test_list_of_zero_cost_models(self, mock_router_with_zero_cost_model): + """Test that a list of zero-cost models returns True.""" + result = _is_model_cost_zero( + model=["on-prem-model"], llm_router=mock_router_with_zero_cost_model + ) + assert result is True + + def test_mixed_cost_models(self, mock_router_with_zero_cost_model): + """Test that a list with mixed cost models returns False.""" + with patch("litellm.get_model_info") as mock_get_model_info: + mock_get_model_info.return_value = { + "input_cost_per_token": 0.0000015, + "output_cost_per_token": 0.000002, + } + result = _is_model_cost_zero( + model=["on-prem-model", "cloud-model"], + llm_router=mock_router_with_zero_cost_model, + ) + assert result is False + + +class TestUserBudgetBypass: + """Tests for user budget bypass with zero-cost models.""" + + @pytest.mark.asyncio + async def test_user_over_budget_with_zero_cost_model_allowed( + self, mock_router_with_zero_cost_model, mock_proxy_logging + ): + """Test that user over budget can still use zero-cost models.""" + user_object = LiteLLM_UserTable( + user_id="test-user", + spend=100.0, + max_budget=50.0, + ) + + request_body = {"model": "on-prem-model"} + + # Should not raise BudgetExceededError + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=user_object, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=mock_router_with_zero_cost_model, + proxy_logging_obj=mock_proxy_logging, + valid_token=UserAPIKeyAuth( + token="test-token", + user_id="test-user", + ), + request=MagicMock(), + skip_budget_checks=True, # This is set by user_api_key_auth for zero-cost models + ) + assert result is True + + @pytest.mark.asyncio + async def test_user_over_budget_with_paid_model_blocked( + self, mock_router_with_zero_cost_model, mock_proxy_logging + ): + """Test that user over budget cannot use paid models.""" + user_object = LiteLLM_UserTable( + user_id="test-user", + spend=100.0, + max_budget=50.0, + ) + + request_body = {"model": "cloud-model"} + + with patch("litellm.get_model_info") as mock_get_model_info: + mock_get_model_info.return_value = { + "input_cost_per_token": 0.0000015, + "output_cost_per_token": 0.000002, + } + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body=request_body, + team_object=None, + user_object=user_object, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=mock_router_with_zero_cost_model, + proxy_logging_obj=mock_proxy_logging, + valid_token=UserAPIKeyAuth( + token="test-token", + user_id="test-user", + ), + request=MagicMock(), + ) + + assert exc_info.value.current_cost == 100.0 + assert exc_info.value.max_budget == 50.0 + assert "test-user" in str(exc_info.value) + + +class TestEndUserBudgetBypass: + """Tests for end user budget bypass with zero-cost models.""" + + @pytest.mark.asyncio + async def test_end_user_over_budget_with_zero_cost_model_allowed( + self, mock_router_with_zero_cost_model, mock_proxy_logging + ): + """Test that end user over budget can still use zero-cost models.""" + end_user_budget = LiteLLM_BudgetTable(max_budget=20.0) + end_user_object = LiteLLM_EndUserTable( + user_id="end-user-123", + spend=50.0, + litellm_budget_table=end_user_budget, + blocked=False, + ) + + request_body = {"model": "on-prem-model", "user": "end-user-123"} + + # In the real flow, skip_budget_checks would be set to True for zero-cost models + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=end_user_object, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=mock_router_with_zero_cost_model, + proxy_logging_obj=mock_proxy_logging, + valid_token=UserAPIKeyAuth( + token="test-token", + ), + request=MagicMock(), + skip_budget_checks=True, # This is set by user_api_key_auth for zero-cost models + ) + assert result is True + + @pytest.mark.asyncio + async def test_end_user_over_budget_with_paid_model_blocked( + self, mock_router_with_zero_cost_model, mock_proxy_logging + ): + """Test that end user over budget cannot use paid models.""" + end_user_budget = LiteLLM_BudgetTable(max_budget=20.0) + end_user_object = LiteLLM_EndUserTable( + user_id="end-user-123", + spend=50.0, + litellm_budget_table=end_user_budget, + blocked=False, + ) + + request_body = {"model": "cloud-model", "user": "end-user-123"} + + with patch("litellm.get_model_info") as mock_get_model_info: + mock_get_model_info.return_value = { + "input_cost_per_token": 0.0000015, + "output_cost_per_token": 0.000002, + } + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=end_user_object, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=mock_router_with_zero_cost_model, + proxy_logging_obj=mock_proxy_logging, + valid_token=UserAPIKeyAuth( + token="test-token", + ), + request=MagicMock(), + ) + + assert exc_info.value.current_cost == 50.0 + assert exc_info.value.max_budget == 20.0 + assert "end-user-123" in str(exc_info.value) + + +class TestTeamBudgetBypass: + """Tests for team budget bypass with zero-cost models.""" + + @pytest.mark.asyncio + async def test_team_over_budget_with_zero_cost_model_allowed( + self, mock_router_with_zero_cost_model, mock_proxy_logging + ): + """Test that team over budget can still use zero-cost models.""" + team_object = LiteLLM_TeamTable( + team_id="test-team", + spend=150.0, + max_budget=100.0, + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + team_id="test-team", + ) + + request_body = {"model": "on-prem-model"} + + # In the real flow, skip_budget_checks would be set to True for zero-cost models + result = await common_checks( + request_body=request_body, + team_object=team_object, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=mock_router_with_zero_cost_model, + proxy_logging_obj=mock_proxy_logging, + valid_token=valid_token, + request=MagicMock(), + skip_budget_checks=True, # This is set by user_api_key_auth for zero-cost models + ) + assert result is True + + @pytest.mark.asyncio + async def test_team_over_budget_with_paid_model_blocked( + self, mock_router_with_zero_cost_model, mock_proxy_logging + ): + """Test that team over budget cannot use paid models.""" + team_object = LiteLLM_TeamTable( + team_id="test-team", + spend=150.0, + max_budget=100.0, + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + team_id="test-team", + ) + + request_body = {"model": "cloud-model"} + + with patch("litellm.get_model_info") as mock_get_model_info: + mock_get_model_info.return_value = { + "input_cost_per_token": 0.0000015, + "output_cost_per_token": 0.000002, + } + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body=request_body, + team_object=team_object, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=mock_router_with_zero_cost_model, + proxy_logging_obj=mock_proxy_logging, + valid_token=valid_token, + request=MagicMock(), + ) + + assert exc_info.value.current_cost == 150.0 + assert exc_info.value.max_budget == 100.0 + assert "test-team" in str(exc_info.value) + + +class TestTeamMemberBudgetBypass: + """Tests for team member budget bypass with zero-cost models.""" + + @pytest.mark.asyncio + async def test_team_member_over_budget_with_zero_cost_model_allowed( + self, mock_router_with_zero_cost_model, mock_proxy_logging + ): + """Test that team member over budget can still use zero-cost models.""" + team_object = LiteLLM_TeamTable( + team_id="test-team", + ) + + user_object = LiteLLM_UserTable( + user_id="test-user", + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="test-user", + team_id="test-team", + ) + + member_budget = LiteLLM_BudgetTable(max_budget=30.0) + team_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=60.0, + litellm_budget_table=member_budget, + ) + + request_body = {"model": "on-prem-model"} + + # Mock get_team_membership + with patch( + "litellm.proxy.auth.auth_checks.get_team_membership" + ) as mock_get_membership: + mock_get_membership.return_value = team_membership + + # In the real flow, skip_budget_checks would be set to True for zero-cost models + result = await common_checks( + request_body=request_body, + team_object=team_object, + user_object=user_object, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=mock_router_with_zero_cost_model, + proxy_logging_obj=mock_proxy_logging, + valid_token=valid_token, + request=MagicMock(), + skip_budget_checks=True, # This is set by user_api_key_auth for zero-cost models + ) + assert result is True + + @pytest.mark.asyncio + async def test_team_member_over_budget_with_paid_model_blocked( + self, mock_router_with_zero_cost_model, mock_proxy_logging + ): + """Test that team member over budget cannot use paid models.""" + team_object = LiteLLM_TeamTable( + team_id="test-team", + ) + + user_object = LiteLLM_UserTable( + user_id="test-user", + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="test-user", + team_id="test-team", + ) + + member_budget = LiteLLM_BudgetTable(max_budget=30.0) + team_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=60.0, + litellm_budget_table=member_budget, + ) + + request_body = {"model": "cloud-model"} + + with patch( + "litellm.proxy.auth.auth_checks.get_team_membership" + ) as mock_get_membership: + mock_get_membership.return_value = team_membership + + with patch("litellm.get_model_info") as mock_get_model_info: + mock_get_model_info.return_value = { + "input_cost_per_token": 0.0000015, + "output_cost_per_token": 0.000002, + } + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body=request_body, + team_object=team_object, + user_object=user_object, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=mock_router_with_zero_cost_model, + proxy_logging_obj=mock_proxy_logging, + valid_token=valid_token, + request=MagicMock(), + ) + + assert exc_info.value.current_cost == 60.0 + assert exc_info.value.max_budget == 30.0 + assert "test-user" in str(exc_info.value) + assert "test-team" in str(exc_info.value) + + +class TestEdgeCases: + """Tests for edge cases and error handling.""" + + def test_model_not_in_router(self, mock_router_with_zero_cost_model): + """Test behavior when model is not found in router.""" + with patch("litellm.get_model_info") as mock_get_model_info: + # Simulate model not found + mock_get_model_info.side_effect = Exception("Model not found") + result = _is_model_cost_zero( + model="nonexistent-model", llm_router=mock_router_with_zero_cost_model + ) + # Should return False (conservative approach) + assert result is False + + @pytest.mark.asyncio + async def test_user_under_budget_with_paid_model_allowed( + self, mock_router_with_zero_cost_model, mock_proxy_logging + ): + """Test that user under budget can use paid models normally.""" + user_object = LiteLLM_UserTable( + user_id="test-user", + spend=30.0, + max_budget=100.0, + ) + + request_body = {"model": "cloud-model"} + + with patch("litellm.get_model_info") as mock_get_model_info: + mock_get_model_info.return_value = { + "input_cost_per_token": 0.0000015, + "output_cost_per_token": 0.000002, + } + # Should not raise BudgetExceededError + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=user_object, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=mock_router_with_zero_cost_model, + proxy_logging_obj=mock_proxy_logging, + valid_token=UserAPIKeyAuth( + token="test-token", + user_id="test-user", + ), + request=MagicMock(), + ) + assert result is True + + @pytest.mark.asyncio + async def test_user_under_budget_with_zero_cost_model_allowed( + self, mock_router_with_zero_cost_model, mock_proxy_logging + ): + """Test that user under budget can use zero-cost models normally.""" + user_object = LiteLLM_UserTable( + user_id="test-user", + spend=30.0, + max_budget=100.0, + ) + + request_body = {"model": "on-prem-model"} + + # Should not raise BudgetExceededError + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=user_object, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=mock_router_with_zero_cost_model, + proxy_logging_obj=mock_proxy_logging, + valid_token=UserAPIKeyAuth( + token="test-token", + user_id="test-user", + ), + request=MagicMock(), + ) + assert result is True