From 53b10e4372dcde9a3bfc88c52e4368bd96758658 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 13 Sep 2023 15:12:47 -0700
Subject: [PATCH] allow updating budget manager cost by passing in text

---
 litellm/__pycache__/main.cpython-311.pyc  | Bin 32927 -> 32927 bytes
 litellm/__pycache__/utils.cpython-311.pyc | Bin 101105 -> 101105 bytes
 litellm/budget_manager.py                 |  17 ++++++++++++++---
 litellm/tests/test_budget_manager.py      |  17 +++++++++++++++++
 pyproject.toml                            |   2 +-
 5 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc
index 45a2b185e8a9a7aa124e40cad18f7630f161435d..63d8dc73ef6e44b85eea47f67255a8010d596f5d 100644
GIT binary patch
delta 21
bcmbQ=$TYu^iEBA8FBbz4ILL40n%V#WJxc{i

delta 21
bcmbQ=$TYu^iEBA8FBbz4u#0Tun%V#WJH7=4

diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc
index 2f5100f88ae0eff15ab0d9c4dc7af5b82ae97849..e12757bfe3bdf3a8bd31ec6319bdc8cd4806720b 100644
GIT binary patch
delta 24
ecmex3mF?qHHm>Eoyj%=G;2__~wUvwUc^3d*V+RQU

delta 24
ecmex3mF?qHHm>Eoyj%=GAT83!wUvwUc^3d(z6RL<

diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py
index 473bb93e2..c73a48820 100644
--- a/litellm/budget_manager.py
+++ b/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)
     
diff --git a/litellm/tests/test_budget_manager.py b/litellm/tests/test_budget_manager.py
index 2fc1a029b..bc5738aa0 100644
--- a/litellm/tests/test_budget_manager.py
+++ b/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()
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index cca9d73fb..5e247bbf8 100644
--- a/pyproject.toml
+++ b/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"