diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index dcb5a7cc183..49060117670 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -100,8 +100,7 @@ class BudgetManager: return self.user_dict[user] def projected_cost(self, model: str, messages: list, user: str): - text: Final = "".join(message["content"] for message in messages) - prompt_tokens: Final = litellm.token_counter(model=model, text=text) + prompt_tokens: Final = litellm.token_counter(model=model, messages=messages) prompt_cost, _ = litellm.cost_per_token(model=model, prompt_tokens=prompt_tokens, completion_tokens=0) current_cost: Final = self.user_dict[user].get("current_cost", 0) projected_cost: Final = prompt_cost + current_cost diff --git a/tests/test_litellm/test_budget_manager.py b/tests/test_litellm/test_budget_manager.py new file mode 100644 index 00000000000..fd6d82b6a62 --- /dev/null +++ b/tests/test_litellm/test_budget_manager.py @@ -0,0 +1,48 @@ +import pytest + +from litellm.budget_manager import BudgetManager + + +@pytest.fixture() +def manager() -> BudgetManager: + bm = BudgetManager(project_name="test", client_type="local") + bm.create_budget(total_budget=10, user="u", duration="daily") + return bm + + +def test_projected_cost_string_content(manager: BudgetManager): + cost = manager.projected_cost( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + user="u", + ) + assert cost >= 0 + + +def test_projected_cost_vision_content(manager: BudgetManager): + cost = manager.projected_cost( + model="gpt-4o-mini", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "hi"}, + {"type": "image_url", "image_url": {"url": "http://x/y.png"}}, + ], + } + ], + user="u", + ) + assert cost >= 0 + + +def test_projected_cost_none_and_missing_content(manager: BudgetManager): + assert ( + manager.projected_cost( + model="gpt-4o-mini", + messages=[{"role": "assistant", "content": None}], + user="u", + ) + >= 0 + ) + assert manager.projected_cost(model="gpt-4o-mini", messages=[{"role": "user"}], user="u") >= 0