mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(passthrough): attribute spend and release budget reservation on router-model /vllm and /azure routes
The /vllm and /azure router-model passthrough branches called llm_router.allm_passthrough_route directly with no request metadata, so the cost callback saw no user_api_key and no user_api_key_budget_reservation. Spend for a budgeted virtual key hit neither the key's spend nor the spend logs, and the reservation minted at auth into the shared Redis counter was never released, drifting the counter up until the key falsely tripped BudgetExceededError. Thread the authenticated key's attribution metadata into both calls via the same builder add_litellm_data_to_request uses, so the cost callback attributes spend and reconciles the reservation. Regression tests cover both branches.
This commit is contained in:
parent
3122600e21
commit
d23069e907
2 changed files with 109 additions and 1 deletions
|
|
@ -9,7 +9,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
|
|||
import json
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, cast
|
||||
|
||||
|
|
@ -106,6 +106,28 @@ def is_passthrough_request_streaming(request_body: object) -> bool:
|
|||
return bool(request_body.get("stream", False))
|
||||
|
||||
|
||||
def get_passthrough_router_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, Any]:
|
||||
"""
|
||||
Build the request metadata carrying key-level spend attribution and the
|
||||
pre-call budget reservation for a router-model passthrough request.
|
||||
|
||||
Router-model passthrough branches call ``allm_passthrough_route`` directly,
|
||||
bypassing ``add_litellm_data_to_request``. Without this metadata the cost
|
||||
callback cannot attribute spend to the calling key and never releases the
|
||||
budget reservation minted at auth time, so the shared spend counter drifts
|
||||
up until the key falsely trips ``BudgetExceededError``.
|
||||
"""
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
|
||||
request_data: Final = {"metadata": {}} # mutable-ok: attribution builder + litellm mutate this dict in place
|
||||
LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
|
||||
data=request_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
_metadata_variable_name="metadata",
|
||||
)
|
||||
return request_data["metadata"]
|
||||
|
||||
|
||||
async def llm_passthrough_factory_proxy_route(
|
||||
custom_llm_provider: str,
|
||||
endpoint: str,
|
||||
|
|
@ -346,6 +368,7 @@ async def vllm_proxy_route(
|
|||
params=None,
|
||||
headers=None,
|
||||
cookies=None,
|
||||
metadata=get_passthrough_router_request_metadata(user_api_key_dict),
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -1475,6 +1498,7 @@ async def azure_proxy_route(
|
|||
params=None,
|
||||
headers=None,
|
||||
cookies=None,
|
||||
metadata=get_passthrough_router_request_metadata(user_api_key_dict),
|
||||
)
|
||||
|
||||
if is_streaming_request:
|
||||
|
|
|
|||
|
|
@ -4055,3 +4055,87 @@ class TestVertexAILiveWebsocketPassthrough:
|
|||
assert "use_in_pass_through" in close_kwargs["reason"]
|
||||
assert "default_vertex_config" in close_kwargs["reason"]
|
||||
assert len(close_kwargs["reason"].encode("utf-8")) <= 123
|
||||
|
||||
|
||||
class TestPassthroughRouterModelBudgetReservation:
|
||||
"""
|
||||
Router-model passthrough on /vllm and /azure must thread the calling key's
|
||||
metadata into ``allm_passthrough_route``. Without ``user_api_key`` the spend
|
||||
is attributed to nobody, and without ``user_api_key_budget_reservation`` the
|
||||
pre-call reservation is never released, so the shared spend counter drifts up
|
||||
until the key falsely trips a 429 BudgetExceededError (LIT-5470).
|
||||
"""
|
||||
|
||||
def _key_with_reservation(self) -> UserAPIKeyAuth:
|
||||
reservation = {
|
||||
"reserved_cost": 0.5,
|
||||
"entries": [{"counter_key": "spend:key:hashed-token", "reserved_cost": 0.5}],
|
||||
}
|
||||
return UserAPIKeyAuth(
|
||||
api_key="hashed-token",
|
||||
user_id="u1",
|
||||
team_id="t1",
|
||||
budget_reservation=reservation,
|
||||
)
|
||||
|
||||
def _request(self) -> MagicMock:
|
||||
request = MagicMock(spec=Request)
|
||||
request.method = "POST"
|
||||
request.headers = {"content-type": "application/json"}
|
||||
request.query_params = {}
|
||||
return request
|
||||
|
||||
def _install_recording_router(self, monkeypatch, body: dict) -> list[dict]:
|
||||
import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
captured: list[dict] = []
|
||||
|
||||
class RecordingRouter:
|
||||
async def allm_passthrough_route(self, **kwargs):
|
||||
captured.append(kwargs)
|
||||
return httpx.Response(200, json={"ok": True})
|
||||
|
||||
async def fake_get_request_body(_request):
|
||||
return body
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", RecordingRouter())
|
||||
monkeypatch.setattr(ep, "get_request_body", fake_get_request_body)
|
||||
monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True)
|
||||
return captured
|
||||
|
||||
def _assert_metadata_carries_attribution(self, captured: list[dict], user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
assert len(captured) == 1, "the router-model branch must dispatch exactly once"
|
||||
metadata = captured[0]["metadata"]
|
||||
assert metadata["user_api_key"] == user_api_key_dict.api_key
|
||||
assert metadata["user_api_key_budget_reservation"] is user_api_key_dict.budget_reservation
|
||||
assert metadata["user_api_key_user_id"] == user_api_key_dict.user_id
|
||||
assert metadata["user_api_key_team_id"] == user_api_key_dict.team_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vllm_router_model_threads_key_metadata(self, monkeypatch):
|
||||
user_api_key_dict = self._key_with_reservation()
|
||||
captured = self._install_recording_router(monkeypatch, {"model": "router-model", "stream": False})
|
||||
|
||||
await vllm_proxy_route(
|
||||
endpoint="/chat/completions",
|
||||
request=self._request(),
|
||||
fastapi_response=MagicMock(spec=Response),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
self._assert_metadata_carries_attribution(captured, user_api_key_dict)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_router_model_threads_key_metadata(self, monkeypatch):
|
||||
user_api_key_dict = self._key_with_reservation()
|
||||
captured = self._install_recording_router(monkeypatch, {"model": "gpt-5", "stream": False})
|
||||
|
||||
await azure_proxy_route(
|
||||
endpoint="openai/deployments/gpt-5/chat/completions",
|
||||
request=self._request(),
|
||||
fastapi_response=MagicMock(spec=Response),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
self._assert_metadata_carries_attribution(captured, user_api_key_dict)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue