diff --git a/docs/my-website/docs/proxy/custom_pricing.md b/docs/my-website/docs/proxy/custom_pricing.md index b5fbd0b6c2e..f6762f5e45c 100644 --- a/docs/my-website/docs/proxy/custom_pricing.md +++ b/docs/my-website/docs/proxy/custom_pricing.md @@ -9,7 +9,6 @@ LiteLLM provides flexible cost tracking and pricing customization for all LLM pr - **Custom Pricing** - Override default model costs or set pricing for custom models - **Cost Per Token** - Track costs based on input/output tokens (most common) - **Cost Per Second** - Track costs based on runtime (e.g., Sagemaker) -- **Zero-Cost Models** - Bypass budget checks for free/on-premises models by setting costs to 0 - **[Provider Discounts](./provider_discounts.md)** - Apply percentage-based discounts to specific providers - **[Provider Margins](./provider_margins.md)** - Add fees/margins to LLM costs for internal billing - **Base Model Mapping** - Ensure accurate cost tracking for Azure deployments @@ -107,51 +106,6 @@ There are other keys you can use to specify costs for different scenarios and mo These keys evolve based on how new models handle multimodality. The latest version can be found at [https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). -## Zero-Cost Models (Bypass Budget Checks) - -**Use Case**: You have on-premises or free models that should be accessible even when users exceed their budget limits. - -**Solution** ✅: Set both `input_cost_per_token` and `output_cost_per_token` to `0` (explicitly) to bypass all budget checks for that model. - -:::info - -When a model is configured with zero cost, LiteLLM will automatically skip ALL budget checks (user, team, team member, end-user, organization, and global proxy budget) for requests to that model. - -**Important**: Both costs must be **explicitly set to 0**. If costs are `null` or undefined, the model will be treated as having cost and budget checks will apply. - -::: - -### Configuration Example - -```yaml -model_list: - # On-premises model - free to use - - model_name: on-prem-llama - litellm_params: - model: ollama/llama3 - api_base: http://localhost:11434 - model_info: - input_cost_per_token: 0 # 👈 Explicitly set to 0 - output_cost_per_token: 0 # 👈 Explicitly set to 0 - - # Paid cloud model - budget checks apply - - model_name: gpt-4 - litellm_params: - model: gpt-4 - api_key: os.environ/OPENAI_API_KEY - # No model_info - uses default pricing from cost map -``` - -### Behavior - -With the above configuration: - -- **User over budget** → Can still use `on-prem-llama` ✅, but blocked from `gpt-4` ❌ -- **Team over budget** → Can still use `on-prem-llama` ✅, but blocked from `gpt-4` ❌ -- **End-user over budget** → Can still use `on-prem-llama` ✅, but blocked from `gpt-4` ❌ - -This ensures your free/on-premises models remain accessible regardless of budget constraints, while paid models are still properly governed. - ## Set 'base_model' for Cost Tracking (e.g. Azure deployments) **Problem**: Azure returns `gpt-4` in the response when `azure/gpt-4-1106-preview` is used. This leads to inaccurate cost tracking diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1879b306253..a741869e5fc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -74,75 +74,6 @@ 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], @@ -155,7 +86,6 @@ 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. @@ -207,66 +137,64 @@ async def common_checks( user_object=user_object, ) - # 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. 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 ( @@ -309,7 +237,6 @@ 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 b0c49182eec..bc0c164a0ad 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -586,21 +586,6 @@ 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, @@ -614,7 +599,6 @@ 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 @@ -1006,22 +990,8 @@ 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 not skip_budget_checks and valid_token.team_member_spend is not None: + if 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}" @@ -1085,47 +1055,46 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 param=abbreviate_api_key(api_key=api_key), ) - 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( + # 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 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. 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 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( @@ -1193,7 +1162,6 @@ 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 deleted file mode 100644 index bc818fc0dca..00000000000 --- a/tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py +++ /dev/null @@ -1,590 +0,0 @@ -""" -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