diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc index a7c59fbcb..d4386ac9c 100644 Binary files a/litellm/__pycache__/main.cpython-311.pyc and b/litellm/__pycache__/main.cpython-311.pyc differ diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc index a12a96977..ea63bbb38 100644 Binary files a/litellm/__pycache__/utils.cpython-311.pyc and b/litellm/__pycache__/utils.cpython-311.pyc differ diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index ddfae14d1..ff7de9b7b 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -1,8 +1,8 @@ -import os, json +import os, json, time import litellm from litellm.utils import ModelResponse import requests, threading -from typing import Optional +from typing import Optional, Union, Literal class BudgetManager: def __init__(self, project_name: str, client_type: str = "local", api_base: Optional[str] = None): @@ -41,8 +41,18 @@ class BudgetManager: else: self.user_dict = response["data"] - def create_budget(self, total_budget: float, user: str): - self.user_dict[user] = {"total_budget": total_budget} + def create_budget(self, total_budget: float, user: str, duration: Literal["daily", "weekly", "monthly", "yearly"], created_at: float = time.time()): + if duration == 'daily': + duration_in_days = 1 + elif duration == 'weekly': + duration_in_days = 7 + elif duration == 'monthly': + duration_in_days = 28 + elif duration == 'yearly': + duration_in_days = 365 + else: + raise ValueError('Invalid duration') + self.user_dict[user] = {"total_budget": total_budget, "duration": duration_in_days, "created_at": created_at, "last_updated_at": created_at} return self.user_dict[user] def projected_cost(self, model: str, messages: list, user: str): @@ -84,6 +94,26 @@ class BudgetManager: self.user_dict[user]["model_cost"] = {} return {"user": self.user_dict[user]} + def reset_on_duration(self, user: str): + # Get current and creation time + last_updated_at = self.user_dict[user]["last_updated_at"] + current_time = time.time() + + # Convert duration from days to seconds + duration_in_seconds = self.user_dict[user]["duration"] * 24 * 60 * 60 + + # Check if duration has elapsed + if current_time - last_updated_at >= duration_in_seconds: + # Reset cost if duration has elapsed and update the creation time + self.reset_cost(user) + self.user_dict[user]["last_updated_at"] = current_time + self._save_data_thread() # Save the data + + def update_budget_all_users(self): + for user in self.get_users(): + if "duration" in self.user_dict[user]: + self.reset_on_duration(user) + def _save_data_thread(self): thread = threading.Thread(target=self.save_data) # [Non-Blocking]: saves data without blocking execution thread.start() diff --git a/litellm/tests/test_budget_manager.py b/litellm/tests/test_budget_manager.py index 1ca3d6dbb..c94e729e6 100644 --- a/litellm/tests/test_budget_manager.py +++ b/litellm/tests/test_budget_manager.py @@ -36,8 +36,6 @@ def test_user_budget_enough(): 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(): try: @@ -78,4 +76,37 @@ def test_get_users(): response = budget_manager.get_users() print(response) except: - pytest.fail(f"An error occurred") \ No newline at end of file + pytest.fail(f"An error occurred") + + +## Scenario 5: Reset budget at the end of duration +def test_reset_on_duration(): + try: + # First, set a short duration budget for a user + user = "123456" + budget_manager.create_budget(total_budget=10, user=user, duration="daily") + + # Use some of the budget + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello!"}] + } + if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user=user): + response = litellm.completion(**data) + print(budget_manager.update_cost(completion_obj=response, user=user)) + + assert budget_manager.get_current_cost(user) > 0, f"Test setup failed: Budget did not decrease after completion" + + # Now, we need to simulate the passing of time. Since we don't want our tests to actually take days, we're going + # to cheat a little -- we'll manually adjust the "created_at" time so it seems like a day has passed. + # In a real-world testing scenario, we might instead use something like the `freezegun` library to mock the system time. + one_day_in_seconds = 24 * 60 * 60 + budget_manager.user_dict[user]["last_updated_at"] -= one_day_in_seconds + + # Now the duration should have expired, so our budget should reset + budget_manager.update_budget_all_users() + + # Make sure the budget was actually reset + assert budget_manager.get_current_cost(user) == 0, "Budget didn't reset after duration expired" + except Exception as e: + pytest.fail(f"An error occurred - {str(e)}") \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 306a0a866..3407318b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "0.1.605" +version = "0.1.606" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT License"