diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 7376c3a402..affffaeea0 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -10,6 +10,7 @@ from fastapi import APIRouter, Depends, HTTPException, status import litellm from litellm._logging import verbose_proxy_logger +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.proxy._types import * from litellm.proxy._types import ProviderBudgetResponse, ProviderBudgetResponseObject from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -2765,16 +2766,23 @@ async def provider_budgets() -> ProviderBudgetResponse: provider_budget_response_dict: Dict[str, ProviderBudgetResponseObject] = {} for _provider, _budget_info in provider_budget_config.items(): - if llm_router.router_budget_logger is None: + router_budget_logger = next( + ( + cb + for cb in (llm_router.optional_callbacks or []) + if isinstance(cb, RouterBudgetLimiting) + ), + None, + ) + if router_budget_logger is None: raise ValueError("No router budget logger found") _provider_spend = ( - await llm_router.router_budget_logger._get_current_provider_spend( + await router_budget_logger._get_current_provider_spend(_provider) or 0.0 + ) + _provider_budget_ttl = ( + await router_budget_logger._get_current_provider_budget_reset_at( _provider ) - or 0.0 - ) - _provider_budget_ttl = await llm_router.router_budget_logger._get_current_provider_budget_reset_at( - _provider ) provider_budget_response_object = ProviderBudgetResponseObject( budget_limit=_budget_info.max_budget, diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index d9ae0665ea..069896e9fb 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -18,7 +18,9 @@ import litellm from litellm.proxy._types import SpendLogsPayload from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.proxy_server import app, prisma_client +from litellm.proxy.spend_tracking import spend_management_endpoints from litellm.router import Router +from litellm.types.utils import BudgetConfig ignored_keys = [ "request_id", @@ -32,6 +34,18 @@ ignored_keys = [ "metadata.cold_storage_object_key", ] +MODEL_LIST = [ + { + "model_name": "azure-gpt-4o", + "litellm_params": { + "model": "azure/gpt-4o-mini", + "mock_response": "Hello, world!", + "tags": ["default"], + "base_model": "gpt-4o-mini", + }, + }, +] + @pytest.fixture def client(): @@ -43,6 +57,19 @@ def add_anthropic_api_key_to_env(monkeypatch): monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-api03-1234567890") +@pytest.fixture +def disable_budget_sync(monkeypatch): + """Disable periodic sync during tests""" + + async def noop(*a, **k): + return None + + monkeypatch.setattr( + "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis", + noop, + ) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_with_user_id(client, monkeypatch): # Mock data for the test @@ -1318,3 +1345,70 @@ async def test_view_spend_tags_no_database(client, monkeypatch): # Check the actual error message structure assert "error" in data assert "Database not connected" in data["error"]["message"] + + +@pytest.mark.asyncio +async def test_provider_budget_under(disable_budget_sync): + """Test that router allows completion when under budget""" + provider_budget_config = { + "azure": BudgetConfig(max_budget=0.01, budget_duration="10d") + } + + router = Router( + enable_pre_call_checks=True, + provider_budget_config=provider_budget_config, + model_list=MODEL_LIST, + ) + + response = await router.acompletion( + model="azure-gpt-4o", + messages=[{"role": "user", "content": "Hello, world!"}], + ) + + assert response is not None + + +@pytest.mark.asyncio +async def test_provider_budget_over(disable_budget_sync): + """Test that router allows completion when over budget""" + provider_budget_config = { + "azure": BudgetConfig(max_budget=-0.01, budget_duration="10d") + } + + router = Router( + num_retries=0, + enable_pre_call_checks=True, + provider_budget_config=provider_budget_config, + model_list=MODEL_LIST, + ) + + with pytest.raises(Exception) as e: + response = await router.acompletion( + model="azure-gpt-4o", + messages=[{"role": "user", "content": "Hello, world!"}], + ) + assert "Exceeded budget for provider" in str(e.value) + + +@pytest.mark.asyncio +async def test_provider_budget_provider_budgets(disable_budget_sync): + """Test that provider_budgets() returns correct values""" + provider = "azure" + max_budget = -0.01 + budget_duration = "10d" + provider_budget_config = { + provider: BudgetConfig(max_budget=max_budget, budget_duration=budget_duration) + } + + router = Router( + num_retries=0, + enable_pre_call_checks=True, + provider_budget_config=provider_budget_config, + model_list=MODEL_LIST, + ) + + with patch("litellm.proxy.proxy_server.llm_router", router): + response = await spend_management_endpoints.provider_budgets() + provider_budget_response = response.providers[provider] + assert provider_budget_response.budget_limit == max_budget + assert provider_budget_response.time_period == budget_duration