mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge 1ce6d571d7 into 44d84360fb
This commit is contained in:
commit
299f5fe6a1
2 changed files with 94 additions and 13 deletions
|
|
@ -222,20 +222,21 @@ class RouterBudgetLimiting(CustomLogger):
|
|||
provider = self._get_llm_provider_for_deployment(deployment)
|
||||
if provider in provider_configs:
|
||||
config = provider_configs[provider]
|
||||
if config.max_budget is None:
|
||||
continue
|
||||
current_spend = spend_map.get(f"provider_spend:{provider}:{config.budget_duration}", 0.0)
|
||||
self._track_provider_remaining_budget_prometheus(
|
||||
provider=provider,
|
||||
spend=current_spend,
|
||||
budget_limit=config.max_budget,
|
||||
)
|
||||
if config.max_budget is not None:
|
||||
current_spend = spend_map.get(f"provider_spend:{provider}:{config.budget_duration}", 0.0)
|
||||
self._track_provider_remaining_budget_prometheus(
|
||||
provider=provider,
|
||||
spend=current_spend,
|
||||
budget_limit=config.max_budget,
|
||||
)
|
||||
|
||||
if config.max_budget and current_spend >= config.max_budget:
|
||||
debug_msg = f"Exceeded budget for provider {provider}: {current_spend} >= {config.max_budget}"
|
||||
deployment_above_budget_info += f"{debug_msg}\n"
|
||||
is_within_budget = False
|
||||
continue
|
||||
if current_spend >= config.max_budget:
|
||||
debug_msg = (
|
||||
f"Exceeded budget for provider {provider}: {current_spend} >= {config.max_budget}"
|
||||
)
|
||||
deployment_above_budget_info += f"{debug_msg}\n"
|
||||
is_within_budget = False
|
||||
continue
|
||||
|
||||
# Check deployment budget
|
||||
if self.deployment_budget_config and is_within_budget:
|
||||
|
|
|
|||
|
|
@ -141,6 +141,86 @@ async def test_async_filter_deployments_resolves_provider_once_per_deployment(
|
|||
assert provider_resolution_calls == len(healthy_deployments)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_filter_deployments_keeps_provider_without_max_budget(
|
||||
disable_budget_sync, monkeypatch
|
||||
):
|
||||
provider_budget = RouterBudgetLimiting(
|
||||
dual_cache=DualCache(),
|
||||
provider_budget_config={
|
||||
"openai": BudgetConfig(budget_duration="1d"),
|
||||
},
|
||||
)
|
||||
|
||||
healthy_deployments = [
|
||||
{
|
||||
"model_name": "gpt-4o-mini",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini"},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-4o-mini",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini"},
|
||||
"model_info": {"id": "deployment-2"},
|
||||
},
|
||||
]
|
||||
|
||||
monkeypatch.setattr(
|
||||
provider_budget,
|
||||
"_get_llm_provider_for_deployment",
|
||||
lambda deployment: "openai",
|
||||
)
|
||||
|
||||
filtered_deployments = await provider_budget.async_filter_deployments(
|
||||
model="gpt-4o-mini",
|
||||
healthy_deployments=healthy_deployments,
|
||||
messages=[],
|
||||
request_kwargs={},
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
||||
assert filtered_deployments == healthy_deployments
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_filter_deployments_filters_provider_over_max_budget(
|
||||
disable_budget_sync, monkeypatch
|
||||
):
|
||||
provider_budget = RouterBudgetLimiting(
|
||||
dual_cache=DualCache(),
|
||||
provider_budget_config={
|
||||
"openai": BudgetConfig(budget_duration="1d", max_budget=100.0),
|
||||
},
|
||||
)
|
||||
|
||||
await provider_budget.dual_cache.async_set_cache(
|
||||
key="provider_spend:openai:1d", value=150.0
|
||||
)
|
||||
|
||||
healthy_deployments = [
|
||||
{
|
||||
"model_name": "gpt-4o-mini",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini"},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
]
|
||||
|
||||
monkeypatch.setattr(
|
||||
provider_budget,
|
||||
"_get_llm_provider_for_deployment",
|
||||
lambda deployment: "openai",
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Exceeded budget for provider openai"):
|
||||
await provider_budget.async_filter_deployments(
|
||||
model="gpt-4o-mini",
|
||||
healthy_deployments=healthy_deployments,
|
||||
messages=[],
|
||||
request_kwargs={},
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_filter_deployments_does_not_recompute_provider_when_resolved_none(
|
||||
disable_budget_sync, monkeypatch
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue