phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 03:34:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add headers to budget manager

Browse source
This commit is contained in:
Jakob 2024-01-18 16:10:45 -08:00
parent e6ec5a2f8c
commit cbe22b2b49
1 changed files with 4 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/budget_manager.py
View file

@ -11,10 +11,12 @@ class BudgetManager:
project_name: str, project_name: str,
client_type: str = "local", client_type: str = "local",
api_base: Optional[str] = None, api_base: Optional[str] = None,
headers: Optional[dict] = None,
): ):
self.client_type = client_type self.client_type = client_type
self.project_name = project_name self.project_name = project_name
self.api_base = api_base or "https://api.litellm.ai" 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 ## load the data or init the initial dictionaries
self.load_data() self.load_data()
@ -43,7 +45,7 @@ class BudgetManager:
url = self.api_base + "/get_budget" url = self.api_base + "/get_budget"
headers = {"Content-Type": "application/json"} headers = {"Content-Type": "application/json"}
data = {"project_name": self.project_name} 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() response = response.json()
if response["status"] == "error": if response["status"] == "error":
self.user_dict = ( self.user_dict = (
@ -201,6 +203,6 @@ class BudgetManager:
url = self.api_base + "/set_budget" url = self.api_base + "/set_budget"
headers = {"Content-Type": "application/json"} headers = {"Content-Type": "application/json"}
data = {"project_name": self.project_name, "user_dict": self.user_dict} 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() response = response.json()
return response return response