scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

Merge pull request #19168 from BerriAI/revert-19027-litellm_add_0_budget_model_bypass

Browse source 
Revert "[Feat] Add support for 0 cost models"
This commit is contained in:
YutaSaito 2026-01-16 10:21:51 +09:00 • committed by GitHub
commit 05a1d15c65
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 87 additions and 828 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
docs/my-website/docs/proxy/custom_pricing.md
View file

@ -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

179
litellm/proxy/auth/auth_checks.py
View file

@ -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

100
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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:

590
tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py
View file

@ -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