minor fix for provider budget

This commit is contained in:
Ishaan Jaff 2024-11-19 20:37:09 -08:00
parent 3c6fe21935
commit 06eb5ecab5

View file

@ -43,6 +43,18 @@ else:
class ProviderBudgetLimiting(CustomLogger): class ProviderBudgetLimiting(CustomLogger):
def __init__(self, router_cache: DualCache, provider_budget_config: dict): def __init__(self, router_cache: DualCache, provider_budget_config: dict):
self.router_cache = router_cache 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 self.provider_budget_config: ProviderBudgetConfigType = provider_budget_config
verbose_router_logger.debug( verbose_router_logger.debug(
f"Initalized Provider budget config: {self.provider_budget_config}" f"Initalized Provider budget config: {self.provider_budget_config}"