From 50eb12ff5b78263bb7841e489e5e6c6e9747a693 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9=20Ferreira?= Date: Mon, 3 Jun 2024 18:12:00 +0200 Subject: [PATCH] test: :white_check_mark: Fix tests with latest changes --- tests/test_costs.py | 4 ++-- tokencost/__init__.py | 1 + tokencost/costs.py | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_costs.py b/tests/test_costs.py index bd3b6a6..ece4b29 100644 --- a/tests/test_costs.py +++ b/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 diff --git a/tokencost/__init__.py b/tokencost/__init__.py index 59e273b..f02126d 100644 --- a/tokencost/__init__.py +++ b/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 diff --git a/tokencost/costs.py b/tokencost/costs.py index bb6d23a..2cf4a1f 100644 --- a/tokencost/costs.py +++ b/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