allow updating budget manager cost by passing in text

litellm/budget_manager.py

@@ -69,18 +69,29 @@ class BudgetManager:
     
     def get_total_budget(self, user: str):
         return self.user_dict[user]["total_budget"]
+    
+    def update_cost(self, user: str, completion_obj: Optional[ModelResponse] = None, model: Optional[str] = None, input_text: Optional[str] = None, output_text: Optional[str] = None):
+        if model and input_text and output_text:
+            prompt_tokens = litellm.token_counter(model=model, messages=[{"role": "user", "content": input_text}])
+            completion_tokens = litellm.token_counter(model=model, messages=[{"role": "user", "content": output_text}])
+            prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar = litellm.cost_per_token(model=model, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens)
+            cost = prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar
+        elif completion_obj: 
+            cost = litellm.completion_cost(completion_response=completion_obj)
+            model = completion_obj['model'] # if this throws an error try, model = completion_obj['model']
+        else:
+            raise ValueError("Either a chat completion object or the text response needs to be passed in. Learn more - https://docs.litellm.ai/docs/budget_manager")
 
-    def update_cost(self, completion_obj: ModelResponse, user: str):
-        cost = litellm.completion_cost(completion_response=completion_obj)
-        model = completion_obj['model'] # if this throws an error try, model = completion_obj['model']
         self.user_dict[user]["current_cost"] = cost + self.user_dict[user].get("current_cost", 0)
         if "model_cost" in self.user_dict[user]:
             self.user_dict[user]["model_cost"][model] = cost + self.user_dict[user]["model_cost"].get(model, 0)
         else:
             self.user_dict[user]["model_cost"] = {model: cost}
+
         self._save_data_thread() # [Non-Blocking] Update persistent storage without blocking execution
         return {"user": self.user_dict[user]}
     
+    
     def get_current_cost(self, user):
         return self.user_dict[user].get("current_cost", 0)
     

litellm/tests/test_budget_manager.py

@@ -110,3 +110,20 @@ def test_reset_on_duration():
         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)}")
+
+## Scenario 6: passing in text: 
+def test_input_text_on_completion():
+    try:
+        user = "12345"
+        budget_manager.create_budget(total_budget=10, user=user, duration="daily")
+
+        input_text = "hello world"
+        output_text = "it's a sunny day in san francisco"
+        model = "gpt-3.5-turbo"
+
+        budget_manager.update_cost(user=user, model=model, input_text=input_text, output_text=output_text)
+        print(budget_manager.get_current_cost(user))
+    except Exception as e:
+        pytest.fail(f"An error occurred - {str(e)}")
+
+test_input_text_on_completion()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.619"
+version = "0.1.620"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"