diff --git a/litellm/router_strategy/provider_budgets.py b/litellm/router_strategy/provider_budgets.py index c1805fea9..7dc3a74e8 100644 --- a/litellm/router_strategy/provider_budgets.py +++ b/litellm/router_strategy/provider_budgets.py @@ -43,6 +43,18 @@ else: class ProviderBudgetLimiting(CustomLogger): def __init__(self, router_cache: DualCache, provider_budget_config: dict): self.router_cache = router_cache + + # cast elements of provider_budget_config to ProviderBudgetInfo + for provider, config in provider_budget_config.items(): + if config is None: + 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"), + ) + self.provider_budget_config: ProviderBudgetConfigType = provider_budget_config verbose_router_logger.debug( f"Initalized Provider budget config: {self.provider_budget_config}"