scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-01 02:02:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 92

fix(agents): preserve paid invocation budgets with free model fields

Browse source
This commit is contained in:
Joshua Valluru 2026-09-27 13:06:32 -07:00
parent 524f819afa
commit 520cea73e4
2 changed files with 32 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -2975,6 +2975,7 @@ async def _run_centralized_common_checks(
request=request,
llm_router=llm_router,
team_id=user_api_key_auth_obj.team_id,
agent_invocation_cost=user_api_key_auth_obj.agent_invocation_cost,
)
# Pin the metadata variable name (litellm_metadata vs metadata) before
@ -3149,7 +3150,10 @@ def _should_skip_budget_checks(
request: Request | None,
llm_router: Any | None,
team_id: str | None = None,
agent_invocation_cost: float | None = None,
) -> bool:
if agent_invocation_cost is not None and agent_invocation_cost > 0:
return False
model: Final = _get_model_from_request_context(
request_data=request_data,
route=route,

28
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -9501,6 +9501,8 @@ async def test_managed_jwt_cannot_be_downgraded_into_virtual_key_mapping(monkeyp
("POST", "/a2a/agent", {"jsonrpc": "2.0", "id": "1", "method": "tasks/get", "params": {"id": "t"}}, False),
("POST", "/a2a/agent", {"jsonrpc": "2.0", "id": "1", "method": "message/send", "params": {}}, True),
("POST", "/a2a/agent", {"jsonrpc": "2.0", "id": "1", "method": "message/stream", "params": {}}, True),
("POST", "/a2a/agent", {"method": "message/send", "model": "free-model", "params": {}}, True),
("POST", "/a2a/agent", {"method": "message/stream", "model": "free-model", "params": {}}, True),
],
)
async def test_human_agent_discovery_does_not_reserve_target_budget_but_send_and_stream_do(
@ -9541,6 +9543,13 @@ async def test_human_agent_discovery_does_not_reserve_target_budget_but_send_and
for name, value in {
**_proxy_attrs_for_centralized_checks(),
"prisma_client": database,
"llm_router": litellm.Router(
model_list=[{
"model_name": "free-model",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
"model_info": {"input_cost_per_token": 0, "output_cost_per_token": 0},
}]
) if body.get("model") else None,
"proxy_logging_obj": MagicMock(post_call_failure_hook=AsyncMock(return_value=None)),
}.items():
monkeypatch.setattr(proxy_server, name, value)
@ -9566,3 +9575,22 @@ async def test_human_agent_discovery_does_not_reserve_target_budget_but_send_and
reserve.assert_awaited_once()
reserved: Final = reserve.call_args.kwargs["valid_token"]
assert reserved is auth and (reserved.billing_agent_policy is not None) is billed, (http_method, body.get("method"))
@pytest.mark.parametrize("invocation_cost,skipped", [(None, True), (0.0, True), (0.25, False)])
def test_free_model_only_waives_budgets_without_a_paid_agent_invocation(
invocation_cost: float | None, skipped: bool
) -> None:
from typing import Final
from litellm.proxy.auth.user_api_key_auth import _should_skip_budget_checks
router: Final = litellm.Router(model_list=[{
"model_name": "free-model",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
"model_info": {"input_cost_per_token": 0, "output_cost_per_token": 0},
}])
assert _should_skip_budget_checks(
request_data={"model": "free-model"}, route="/chat/completions", request=None,
llm_router=router, agent_invocation_cost=invocation_cost,
) is skipped