test: ✅ Fix tests with latest changes

tests/test_costs.py

@@ -151,7 +151,7 @@ def test_count_string_invalid_model():
 def test_calculate_prompt_cost(prompt, model, expected_output):
     """Test that the cost calculation is correct."""
 
-    cost, _ = calculate_prompt_cost(prompt, model)
+    cost = calculate_prompt_cost(prompt, model)
     assert cost == expected_output
 
 
@@ -185,7 +185,7 @@ def test_invalid_prompt_format():
 def test_calculate_completion_cost(prompt, model, expected_output):
     """Test that the completion cost calculation is correct."""
 
-    cost, _ = calculate_completion_cost(prompt, model)
+    cost = calculate_completion_cost(prompt, model)
     assert cost == expected_output
 
 

tokencost/__init__.py

@@ -3,5 +3,6 @@ from .costs import (
     count_string_tokens,
     calculate_completion_cost,
     calculate_prompt_cost,
+    calculate_all_costs_and_tokens,
 )
 from .constants import TOKEN_COSTS_STATIC, TOKEN_COSTS, update_token_costs

tokencost/costs.py

@@ -3,7 +3,7 @@ Costs dictionary and utility tool for counting tokens
 """
 
 import tiktoken
-from typing import Tuple, Union, List, Dict
+from typing import Union, List, Dict
 from .constants import TOKEN_COSTS
 from decimal import Decimal
 import logging