This commit is contained in:
devin-ai-integration[bot] 2026-08-27 18:14:03 -05:00 committed by GitHub
commit 299f5fe6a1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 94 additions and 13 deletions

View file

@ -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:

View file

@ -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