From 06eb5ecab591149db2c2336155b0e08e04705d71 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 19 Nov 2024 20:37:09 -0800
Subject: [PATCH] minor fix for provider budget

---
 litellm/router_strategy/provider_budgets.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

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}"