mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(proxy): use current spend for virtual key soft budget alerts
This commit is contained in:
Tanishq
parent
dddbaa803d
commit
8c0fe877b9
3 changed files with 174 additions and 5 deletions
|
|
@ -6801,16 +6801,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,
|
||||
|
|
|
|||
|
|
@ -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: `<clock>`", message)
|
||||
assert {**payload, "text": normalized} == {
|
||||
"text": (
|
||||
"Alert type: `budget_alerts`\nLevel: `High`\nTimestamp: `<clock>`\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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue