diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc index 45a2b185e..63d8dc73e 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 2f5100f88..e12757bfe 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 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"