fix: Update OpenAI pricing per https://openai.com/api/pricing/ (#110)

Co-authored-by: Lifei Zhou <lifei@squareup.com>
This commit is contained in:
Ben Walding
2024-10-03 21:56:08 +00:00
committed by GitHub
parent 9c0501b3f9
commit 9e35c6370e
2 changed files with 14 additions and 4 deletions
+12 -3
View File
@@ -1,13 +1,22 @@
from unittest.mock import patch
import pytest
from goose.utils._cost_calculator import _calculate_cost, get_total_cost_message
from exchange.providers.base import Usage
def test_calculate_cost():
@pytest.fixture
def mock_prices():
prices = {"gpt-4o": (5.00, 15.00), "gpt-4o-mini": (0.150, 0.600)}
with patch("goose.utils._cost_calculator.PRICES", prices) as mock_prices:
yield mock_prices
def test_calculate_cost(mock_prices):
cost = _calculate_cost("gpt-4o", Usage(input_tokens=10000, output_tokens=600, total_tokens=10600))
assert cost == 0.059
def test_get_total_cost_message():
def test_get_total_cost_message(mock_prices):
message = get_total_cost_message(
{
"gpt-4o": Usage(input_tokens=10000, output_tokens=600, total_tokens=10600),
@@ -22,7 +31,7 @@ def test_get_total_cost_message():
assert message == expected_message
def test_get_total_cost_message_with_non_available_pricing():
def test_get_total_cost_message_with_non_available_pricing(mock_prices):
message = get_total_cost_message(
{
"non_pricing_model": Usage(input_tokens=10000, output_tokens=600, total_tokens=10600),