diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index 4a3bb2cae2..0364741979 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -11,10 +11,12 @@ class BudgetManager: project_name: str, client_type: str = "local", api_base: Optional[str] = None, + headers: Optional[dict] = None, ): self.client_type = client_type self.project_name = project_name self.api_base = api_base or "https://api.litellm.ai" + self.headers = headers or {'Content-Type': 'application/json'} ## load the data or init the initial dictionaries self.load_data() @@ -43,7 +45,7 @@ class BudgetManager: url = self.api_base + "/get_budget" headers = {"Content-Type": "application/json"} data = {"project_name": self.project_name} - response = requests.post(url, headers=headers, json=data) + response = requests.post(url, headers=self.headers, json=data) response = response.json() if response["status"] == "error": self.user_dict = ( @@ -201,6 +203,6 @@ class BudgetManager: url = self.api_base + "/set_budget" headers = {"Content-Type": "application/json"} data = {"project_name": self.project_name, "user_dict": self.user_dict} - response = requests.post(url, headers=headers, json=data) + response = requests.post(url, headers=self.headers, json=data) response = response.json() return response