diff --git a/litellm/router_strategy/provider_budgets.py b/litellm/router_strategy/provider_budgets.py index 6ccbe78d6..37a6e72c2 100644 --- a/litellm/router_strategy/provider_budgets.py +++ b/litellm/router_strategy/provider_budgets.py @@ -51,10 +51,12 @@ class ProviderBudgetLimiting(CustomLogger): raise ValueError( f"No budget config found for provider {provider}, provider_budget_config: {provider_budget_config}" ) - provider_budget_config[provider] = ProviderBudgetInfo( - budget_limit=config.get("budget_limit"), - time_period=config.get("time_period"), - ) + + if not isinstance(config, ProviderBudgetInfo): + provider_budget_config[provider] = ProviderBudgetInfo( + budget_limit=config.get("budget_limit"), + time_period=config.get("time_period"), + ) self.provider_budget_config: ProviderBudgetConfigType = provider_budget_config verbose_router_logger.debug( diff --git a/tests/local_testing/test_provider_budgets.py b/tests/local_testing/test_provider_budgets.py index 5e685cae6..7d1cc4fd6 100644 --- a/tests/local_testing/test_provider_budgets.py +++ b/tests/local_testing/test_provider_budgets.py @@ -93,8 +93,14 @@ async def test_provider_budgets_e2e_test_expect_to_fail(): - first request passes, all subsequent requests fail """ - provider_budget_config: ProviderBudgetConfigType = { - "anthropic": ProviderBudgetInfo(time_period="1d", budget_limit=0.000000000001), + + # Note: We intentionally use a dictionary with string keys for budget_limit and time_period + # we want to test that the router can handle type conversion, since the proxy config yaml passes these values as a dictionary + provider_budget_config = { + "anthropic": { + "budget_limit": 0.000000000001, + "time_period": "1d", + } } router = Router(