feat: track cost and token usage in log file (#80)

tests/cli/test_session.py

@@ -167,3 +167,14 @@ def test_generate_session_name(create_session_with_mock_configs):
         session.generate_session_name()
 
         assert session.name == SPECIFIED_SESSION_NAME
+
+
+def test_log_log_cost(create_session_with_mock_configs):
+    session = create_session_with_mock_configs()
+    mock_logger = MagicMock()
+    cost_message = "You have used 100 tokens"
+    with patch("exchange.Exchange.get_token_usage", return_value={}), patch(
+        "goose.cli.session.get_total_cost_message", return_value=cost_message
+    ), patch("goose.cli.session.get_logger", return_value=mock_logger):
+        session._log_cost()
+        mock_logger.info.assert_called_once_with(cost_message)

tests/utils/test_cost_calculator.py

@@ -0,0 +1,38 @@
+from goose.utils._cost_calculator import _calculate_cost, get_total_cost_message
+from exchange.providers.base import Usage
+
+
+def test_calculate_cost():
+    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():
+    message = get_total_cost_message(
+        {
+            "gpt-4o": Usage(input_tokens=10000, output_tokens=600, total_tokens=10600),
+            "gpt-4o-mini": Usage(input_tokens=3000000, output_tokens=4000000, total_tokens=7000000),
+        }
+    )
+    expected_message = (
+        "Cost for model gpt-4o Usage(input_tokens=10000, output_tokens=600, total_tokens=10600): $0.06\n"
+        + "Cost for model gpt-4o-mini Usage(input_tokens=3000000, output_tokens=4000000, total_tokens=7000000)"
+        + ": $2.85\nTotal cost: $2.91"
+    )
+    assert message == expected_message
+
+
+def test_get_total_cost_message_with_non_available_pricing():
+    message = get_total_cost_message(
+        {
+            "non_pricing_model": Usage(input_tokens=10000, output_tokens=600, total_tokens=10600),
+            "gpt-4o-mini": Usage(input_tokens=3000000, output_tokens=4000000, total_tokens=7000000),
+        }
+    )
+    expected_message = (
+        "Cost for model non_pricing_model Usage(input_tokens=10000, output_tokens=600, total_tokens=10600): "
+        + "Not available\n"
+        + "Cost for model gpt-4o-mini Usage(input_tokens=3000000, output_tokens=4000000, total_tokens=7000000)"
+        + ": $2.85\nTotal cost: $2.85"
+    )
+    assert message == expected_message