From 3fa749fbf6bf2807d0c0d1a973732aaa779c8052 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 11 Sep 2023 15:22:52 -0700 Subject: [PATCH] updates to budget manager naming schema --- docs/my-website/docs/budget_manager.md | 17 +++++++++-------- litellm/budget_manager.py | 18 +++++++++++------- litellm/tests/test_budget_manager.py | 9 ++++++--- pyproject.toml | 2 +- 4 files changed, 27 insertions(+), 19 deletions(-) diff --git a/docs/my-website/docs/budget_manager.md b/docs/my-website/docs/budget_manager.md index d247e06bcb..f7ae84dbe3 100644 --- a/docs/my-website/docs/budget_manager.md +++ b/docs/my-website/docs/budget_manager.md @@ -11,8 +11,9 @@ budget_manager = BudgetManager(project_name="test_project") user = "1234" -# create a budget for a user -budget_manager.create_budget(total_budget=10, user=user) +# create a budget if new user user +if not budget_manager.is_valid_user(user): + budget_manager.create_budget(total_budget=10, user=user) # check if a given call can be made if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user): @@ -40,7 +41,7 @@ budget_manager.get_model_cost(user=user) # {"gpt-3.5-turbo-0613": 7.3e-05} ### save budget to disk -When you call `save_data()` it will check for the self.type (by default this is set to client), and save the dictionary to a local `user_cost.json` file. +When you call `save_data()` it will check for the self.client_type (by default this is set to local), and save the dictionary to a local `user_cost.json` file. ```python # ... @@ -49,20 +50,20 @@ budget_manager.save_data() # 👈 save to user_cost.json() [**Implementation Code**](https://github.com/BerriAI/litellm/blob/817798c692207569a17c26186d10541aa83f04e7/litellm/budget_manager.py#L83) -### save budget to db (LiteLLM) +### save budget to hosted db (LiteLLM) -Set the BudgetManager type to `client`. +Set the BudgetManager type to `client_type`. ```python -budget_manager = BudgetManager(project_name="test_project", type="client") +budget_manager = BudgetManager(project_name="test_project", client_type="hosted") # ... budget_manager.save_data() # 👈 saved to hosted db ``` [**Implementation Code**](https://github.com/BerriAI/litellm/blob/817798c692207569a17c26186d10541aa83f04e7/litellm/budget_manager.py#L11) -### save budget to db (self-hosted) +### save budget to hosted db (self-hosted) -Set the BudgetManager type to `client`. Overwrite the api_base +Set the BudgetManager type to `client_type`. Overwrite the api_base ```python budget_manager = BudgetManager(project_name="test_project", type="client", api_base="your_custom_api") # ... diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index 63c81d7097..9eb8ffd037 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -5,8 +5,8 @@ import requests from typing import Optional class BudgetManager: - def __init__(self, project_name: str, type: str = "local", api_base: Optional[str] = None): - self.type = type + def __init__(self, project_name: str, client_type: str = "local", api_base: Optional[str] = None): + self.client_type = client_type self.project_name = project_name self.api_base = api_base or "https://api.litellm.ai" ## load the data or init the initial dictionaries @@ -17,7 +17,7 @@ class BudgetManager: print(print_statement) def load_data(self): - if self.type == "local": + if self.client_type == "local": # Check if user dict file exists if os.path.isfile("user_cost.json"): # Load the user dict @@ -26,7 +26,8 @@ class BudgetManager: else: self.print_verbose("User Dictionary not found!") self.user_dict = {} - elif self.type == "client": + self.print_verbose(f"user dict from local: {self.user_dict}") + elif self.client_type == "hosted": # Load the user_dict from hosted db url = self.api_base + "/get_budget" headers = {'Content-Type': 'application/json'} @@ -70,7 +71,10 @@ class BudgetManager: def get_model_cost(self, user): return self.user_dict[user].get("model_cost", 0) - + + def is_valid_user(self, user: str) -> bool: + return user in self.user_dict + def get_users(self): return list(self.user_dict.keys()) @@ -80,14 +84,14 @@ class BudgetManager: return {"user": self.user_dict[user]} def save_data(self): - if self.type == "local": + if self.client_type == "local": import json # save the user dict with open("user_cost.json", 'w') as json_file: json.dump(self.user_dict, json_file, indent=4) # Indent for pretty formatting return {"status": "success"} - elif self.type == "client": + elif self.client_type == "hosted": url = self.api_base + "/set_budget" headers = {'Content-Type': 'application/json'} data = { diff --git a/litellm/tests/test_budget_manager.py b/litellm/tests/test_budget_manager.py index 79af9a2f06..1ca3d6dbbd 100644 --- a/litellm/tests/test_budget_manager.py +++ b/litellm/tests/test_budget_manager.py @@ -9,9 +9,10 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import litellm +litellm.set_verbose = True from litellm import BudgetManager, completion -budget_manager = BudgetManager(project_name="test_project") +budget_manager = BudgetManager(project_name="test_project", client_type="hosted") ## Scenario 1: User budget enough to make call def test_user_budget_enough(): @@ -32,8 +33,10 @@ def test_user_budget_enough(): response = "Sorry - no budget!" print(f"response: {response}") - except: - pytest.fail(f"An error occurred") + except Exception as e: + pytest.fail(f"An error occurred - {str(e)}") + +test_user_budget_enough() ## Scenario 2: User budget not enough to make call def test_user_budget_not_enough(): diff --git a/pyproject.toml b/pyproject.toml index 0072a9e1d5..ac9723d0b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "0.1.595" +version = "0.1.596" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT License"