add reset budget on duration end to budget manager

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()

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") 
+        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)}")

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"