diff --git a/tests/test_costs.py b/tests/test_costs.py index 96f5b86c..724214d0 100644 --- a/tests/test_costs.py +++ b/tests/test_costs.py @@ -271,3 +271,8 @@ def test_calculate_cached_tokens_cost(): # Assert that the costs match assert actual_cost == expected_cost assert actual_cost > 0, "Cache token cost should be greater than zero" + + +def test_calculate_cost_by_tokens_rejects_negative_token_counts(): + with pytest.raises(ValueError, match="num_tokens must be non-negative"): + calculate_cost_by_tokens(-1, "gpt-4o", "input") diff --git a/tokencost/costs.py b/tokencost/costs.py index cc4680ff..7e276372 100644 --- a/tokencost/costs.py +++ b/tokencost/costs.py @@ -317,6 +317,9 @@ def calculate_cost_by_tokens(num_tokens: int, model: str, token_type: TokenType) Returns: Decimal: The calculated cost in USD. """ + if num_tokens < 0: + raise ValueError("num_tokens must be non-negative") + model = _normalize_model_for_pricing(model) if model not in TOKEN_COSTS: raise KeyError(