diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 448643ef59b..21ff2c6efc5 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -6824,16 +6824,25 @@ async def virtual_key_soft_budget_check( """ - if valid_token.soft_budget and valid_token.spend >= valid_token.soft_budget: + if not valid_token.soft_budget: + return + + from litellm.proxy.proxy_server import get_current_spend + + spend: Final = await get_current_spend( + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, + ) + if spend >= valid_token.soft_budget: verbose_proxy_logger.debug( "Crossed Soft Budget for token %s, spend %s, soft_budget %s", valid_token.token, - valid_token.spend, + spend, valid_token.soft_budget, ) call_info: Final = CallInfo( token=valid_token.token, - spend=valid_token.spend, + spend=spend, max_budget=valid_token.max_budget, soft_budget=valid_token.soft_budget, user_id=valid_token.user_id, diff --git a/tests/integration/spend/test_key_budget_lockout.py b/tests/integration/spend/test_key_budget_lockout.py index a3d1d3d0a78..f18b67ee7fa 100644 --- a/tests/integration/spend/test_key_budget_lockout.py +++ b/tests/integration/spend/test_key_budget_lockout.py @@ -1,5 +1,6 @@ import json import os +import re import uuid from hashlib import sha256 from pathlib import Path @@ -7,9 +8,10 @@ from typing import Final import httpx import pytest -from integration._support.client import Gateway, eventually, object_value +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import OwnedState, owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, JsonValue, TypeAdapter @@ -201,6 +203,73 @@ def test_an_exhausted_key_is_refused_inference_but_can_still_read_its_own_info(g assert own["max_budget"] == 0.06 +def test_soft_budget_alert_reads_spend_from_another_proxy(gateway: Gateway, tmp_path: Path) -> None: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "general_settings": { + **object_value(base["general_settings"]), + "alerting": ["slack"], + "alert_types": ["budget_alerts"], + }, + } + path: Final = tmp_path / "soft-budget-alerting.yaml" + path.write_text(yaml.safe_dump(config)) + email: Final = "budget-owner@example.test" + alias: Final = f"soft-budget-{uuid.uuid4().hex}" + with gateway.scenario() as scenario, wire_server(lambda _: Reply(body=b"ok", content_type="text/plain")) as alerts: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + user_id: Final = scenario.user(user_email=email) + key: Final = scenario.key(user_id=user_id, models=[model], key_alias=alias, soft_budget=0.1) + digest: Final = sha256(key.encode()).hexdigest() + overrides: Final = { + "SLACK_WEBHOOK_URL": f"{alerts.url}/alerts", + "PROXY_BASE_URL": "http://proxy.integration.test", + } + with owned_proxy(gateway, tmp_path, overrides, config=path, state=OwnedState.SHARED_WITH_WORKER) as candidate: + first: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "warm the alerting proxy"}]}, + key=key, + ) + assert first.status_code == 200, first.text + assert float(first.headers["x-litellm-response-cost"]) == pytest.approx(20 * 0.001 + 20 * 0.002) + eventually(lambda: _spend(gateway, digest), lambda spend: spend == pytest.approx(0.06), seconds=70) + second: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "cross the budget on the other proxy"}]}, + key=key, + ) + assert second.status_code == 200, second.text + eventually(lambda: _spend(gateway, digest), lambda spend: spend == pytest.approx(0.12), seconds=70) + assert alerts.received.qsize() == 0 + third: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "check the shared spend"}]}, + key=key, + ) + assert third.status_code == 200, third.text + eventually(alerts.received.qsize, lambda count: count >= 1, seconds=40) + requests: Final = alerts.drain() + assert len(requests) == 1, requests + assert (requests[0].method, requests[0].target) == ("POST", "/alerts") + payload: Final = JSON_OBJECT.validate_json(requests[0].body) + message: Final = string_value(payload["text"]) + normalized: Final = re.sub(r"Timestamp: `\d{2}:\d{2}:\d{2}`", "Timestamp: ``", message) + assert {**payload, "text": normalized} == { + "text": ( + "Alert type: `budget_alerts`\nLevel: `High`\nTimestamp: ``\n\n" + "Message: Soft Budget Crossed: Total Soft Budget:`0.1`\n\n" + f"*spend:* `0.12`\n*soft_budget:* `0.1`\n*user_id:* `{user_id}`\n" + f"*user_email:* `{email}`\n*key_alias:* `{alias}`\n*event_group:* `key`\n" + "\n\nProxy URL: `http://proxy.integration.test`" + ) + } + + def _bounded_chat(gateway: Gateway, model: str, key: str, content: str | None = None) -> httpx.Response: return gateway.request( "POST", diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e1afa81d372..d3ed39b498d 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -48,6 +48,8 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.utils import ProxyLogging from litellm.proxy.utils import CallInfo +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER @@ -1063,6 +1065,95 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert): ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" +@pytest.mark.parametrize( + "cached_spend, counter_spend, soft_budget, expected_spend", + [ + (0.0, 100.0, 50.0, 100.0), + (0.0, 50.0, 50.0, 50.0), + (100.0, 25.0, 50.0, None), + (100.0, 0.0, 50.0, None), + (100.0, None, 50.0, 100.0), + (0.0, 100.0, None, None), + (0.0, 100.0, 0.0, None), + ], +) +@pytest.mark.asyncio +async def test_soft_budget_alert_uses_current_spend_without_blocking_requests( + monkeypatch: pytest.MonkeyPatch, + cached_spend: float, + counter_spend: float | None, + soft_budget: float | None, + expected_spend: float | None, +) -> None: + from litellm.proxy import proxy_server + + counter: Final = DualCache() + monkeypatch.setattr(proxy_server, "spend_counter_cache", counter) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setenv("WEBHOOK_URL", "https://alerts.example.test/budget") + if counter_spend is not None: + counter.in_memory_cache.set_cache(key="spend:key:soft-budget-key", value=counter_spend) + received: Final[list[httpx.Request]] = [] + + def receive(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response(200, text="ok") + + handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(receive)) + logging: Final = ProxyLogging(UserApiKeyCache()) + logging.alerting = ["webhook"] + logging.slack_alerting_instance = SlackAlerting(alerting=["webhook"], async_http_handler=handler) + token: Final = UserAPIKeyAuth( + token="soft-budget-key", + spend=cached_spend, + soft_budget=soft_budget, + user_id="budget-user", + org_id="budget-org", + key_alias="budget-key", + ) + tasks_before: Final = asyncio.all_tasks() + try: + await virtual_key_soft_budget_check( + token, logging, LiteLLM_UserTable(user_id="budget-user", user_email="owner@example.test") + ) + await asyncio.gather(*(asyncio.all_tasks() - tasks_before)) + expected: Final = ( + [] + if expected_spend is None + else [ + { + "spend": expected_spend, + "max_budget": None, + "soft_budget": soft_budget, + "token": "soft-budget-key", + "customer_id": None, + "user_id": "budget-user", + "team_id": None, + "team_alias": None, + "organization_id": "budget-org", + "user_email": "owner@example.test", + "key_alias": "budget-key", + "projected_exceeded_date": None, + "projected_spend": None, + "event_group": "key", + "alert_emails": None, + "max_budget_alert_emails": None, + "event": "soft_budget_crossed", + "event_message": f"Soft Budget Crossed: Total Soft Budget:`{soft_budget}`", + } + ] + ) + assert [json.loads(request.content) for request in received] == expected + assert token.spend == cached_spend + finally: + pending: Final = asyncio.all_tasks() - tasks_before + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + await handler.close() + + @pytest.mark.parametrize( "spend, soft_budget, expect_alert, metadata, expected_alert_emails", [ diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 092805a9cac..58f3640863e 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -642,80 +642,102 @@ async def test_auth_with_form_data_and_model(): assert response.models == ["gpt-4"], "Model from virtual key should be preserved" +@pytest.mark.parametrize("cached_spend", [0.0, 15.0]) @pytest.mark.asyncio -async def test_soft_budget_alert(): - """ - Test that when a token's spend exceeds soft_budget, it triggers a budget alert but allows the request - """ +async def test_soft_budget_alert(monkeypatch: pytest.MonkeyPatch, cached_spend: float) -> None: import asyncio import time + import httpx from fastapi import Request - from starlette.datastructures import URL + from litellm.caching.dual_cache import DualCache + from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy import proxy_server from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - from litellm.proxy.proxy_server import hash_token, user_api_key_cache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import PrismaClient, ProxyLogging - # Setup - user_key = "sk-98765" - soft_budget = 10 - current_spend = 15 # Spend exceeds soft budget - - # Create a valid token with soft budget - valid_token = UserAPIKeyAuth( - token=hash_token(user_key), - soft_budget=soft_budget, - spend=current_spend, - last_refreshed_at=time.time(), + monkeypatch.setattr(time, "time", lambda: 1_700_000_000.0) + monkeypatch.setenv("WEBHOOK_URL", "https://alerts.example.test/budget") + user_key: Final = "sk-soft-budget-auth" + token_hash: Final = proxy_server.hash_token(user_key) + key_cache: Final = UserApiKeyCache() + counter: Final = DualCache() + key_cache.set_cache( + key=token_hash, + value=UserAPIKeyAuth( + token=token_hash, + soft_budget=10.0, + spend=cached_spend, + last_refreshed_at=1_700_000_000.0, + ), ) + counter.in_memory_cache.set_cache(key=f"spend:key:{token_hash}", value=15.0) + received: Final[list[httpx.Request]] = [] - # Store in cache - user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token) + def receive_alert(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response(200, text="ok") - # Mock proxy server settings - setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - setattr(litellm.proxy.proxy_server, "prisma_client", AsyncMock()) + async def receive_body() -> Message: + return {"type": "http.request", "body": b"{}", "more_body": False} - # Create request - request = Request(scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}) - request._url = URL(url="/chat/completions") - - # Track if budget_alerts was called - alert_called = False - original_budget_alerts = litellm.proxy.proxy_server.proxy_logging_obj.budget_alerts - - async def mock_budget_alerts(*args, **kwargs): - nonlocal alert_called - if kwargs.get("type") == "soft_budget": - alert_called = True - return await original_budget_alerts(*args, **kwargs) - - # Patch the budget_alerts method - setattr( - litellm.proxy.proxy_server.proxy_logging_obj, - "budget_alerts", - mock_budget_alerts, + handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(receive_alert)) + proxy_logging: Final = ProxyLogging(key_cache) + proxy_logging.alerting = ["webhook"] + proxy_logging.slack_alerting_instance = SlackAlerting(alerting=["webhook"], async_http_handler=handler) + monkeypatch.setattr(proxy_server, "user_api_key_cache", key_cache) + monkeypatch.setattr(proxy_server, "spend_counter_cache", counter) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging) + monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr( + proxy_server, + "prisma_client", + PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging), ) - + request: Final = Request( + scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}, + receive=receive_body, + ) + tasks_before: Final = asyncio.all_tasks() + proxy_logging.slack_alerting_instance.update_values(alerting=["webhook"]) + background_tasks: Final = asyncio.all_tasks() try: - # Call user_api_key_auth - response = await user_api_key_auth(request=request, api_key="Bearer " + user_key) - - # Assert the request was allowed (no exception raised) - assert response is not None - # Assert the alert was triggered - await asyncio.sleep(3) - assert alert_called == True, "Soft budget alert should have been triggered" - + response: Final = await user_api_key_auth(request=request, api_key="Bearer " + user_key) + await asyncio.gather(*(asyncio.all_tasks() - background_tasks)) + assert (response.token, response.spend, response.soft_budget) == (token_hash, cached_spend, 10.0) + assert [json.loads(alert.content) for alert in received] == [ + { + "spend": 15.0, + "max_budget": None, + "soft_budget": 10.0, + "token": token_hash, + "customer_id": None, + "user_id": None, + "team_id": None, + "team_alias": None, + "organization_id": None, + "user_email": None, + "key_alias": None, + "projected_exceeded_date": None, + "projected_spend": None, + "event_group": "key", + "alert_emails": None, + "max_budget_alert_emails": None, + "event": "soft_budget_crossed", + "event_message": "Soft Budget Crossed: Total Soft Budget:`10.0`", + } + ] finally: - # Restore original budget_alerts - setattr( - litellm.proxy.proxy_server.proxy_logging_obj, - "budget_alerts", - original_budget_alerts, - ) + pending: Final = asyncio.all_tasks() - tasks_before + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + await handler.close() def test_is_allowed_route():