From 9e35c6370ea22488eb068a248de5c0670ef9afbb Mon Sep 17 00:00:00 2001 From: Ben Walding Date: Thu, 3 Oct 2024 21:56:08 +0000 Subject: [PATCH] fix: Update OpenAI pricing per https://openai.com/api/pricing/ (#110) Co-authored-by: Lifei Zhou --- src/goose/utils/_cost_calculator.py | 3 ++- tests/utils/test_cost_calculator.py | 15 ++++++++++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/goose/utils/_cost_calculator.py b/src/goose/utils/_cost_calculator.py index ae2f1379..5cdf6950 100644 --- a/src/goose/utils/_cost_calculator.py +++ b/src/goose/utils/_cost_calculator.py @@ -2,8 +2,9 @@ from typing import Optional from exchange.providers.base import Usage PRICES = { - "gpt-4o": (5.00, 15.00), + "gpt-4o": (2.50, 10.00), "gpt-4o-2024-08-06": (2.50, 10.00), + "gpt-4o-2024-05-13": (5.00, 15.00), "gpt-4o-mini": (0.150, 0.600), "gpt-4o-mini-2024-07-18": (0.150, 0.600), "o1-preview": (15.00, 60.00), diff --git a/tests/utils/test_cost_calculator.py b/tests/utils/test_cost_calculator.py index 6a28509e..edde1daa 100644 --- a/tests/utils/test_cost_calculator.py +++ b/tests/utils/test_cost_calculator.py @@ -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),