feat: track cost and token usage in log file

pyproject.toml

@@ -8,7 +8,7 @@ dependencies = [
     "attrs>=23.2.0",
     "rich>=13.7.1",
     "ruamel-yaml>=0.18.6",
-    "ai-exchange>=0.9.0",
+    "ai-exchange",
     "click>=8.1.7",
     "prompt-toolkit>=3.0.47",
 ]
@@ -63,4 +63,5 @@ dev-dependencies = [
     "mkdocs-callouts>=1.14.0",
 ]
 
-
+[tool.uv.sources]
+ai-exchange = { path = "/Users/lifei/Development/exchange" }

src/goose/_logger.py

@@ -0,0 +1,20 @@
+import logging
+from pathlib import Path
+
+_LOGGER_NAME = "goose"
+_LOGGER_FILE_NAME = "goose.log"
+
+
+def setup_logging(log_file_directory: Path, level: int = logging.INFO) -> None:
+    logger = logging.getLogger(_LOGGER_NAME)
+    logger.setLevel(level)
+    log_file_directory.mkdir(parents=True, exist_ok=True)
+    file_handler = logging.FileHandler(log_file_directory / _LOGGER_FILE_NAME)
+    logger.addHandler(file_handler)
+    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    file_handler.setFormatter(formatter)
+    return logger
+
+
+def get_logger() -> logging.Logger:
+    return logging.getLogger(_LOGGER_NAME)

src/goose/cli/config.py

@@ -17,6 +17,7 @@
 PROFILES_CONFIG_PATH = GOOSE_GLOBAL_PATH.joinpath("profiles.yaml")
 SESSIONS_PATH = GOOSE_GLOBAL_PATH.joinpath("sessions")
 SESSION_FILE_SUFFIX = ".jsonl"
+LOG_PATH = GOOSE_GLOBAL_PATH.joinpath("logs")
 
 
 @cache

src/goose/cli/session.py

@@ -12,16 +12,13 @@
 from rich.status import Status
 
 from goose.build import build_exchange
-from goose.cli.config import (
-    default_profiles,
-    ensure_config,
-    read_config,
-    session_path,
-)
+from goose.cli.config import default_profiles, ensure_config, read_config, session_path, LOG_PATH
+from goose._logger import get_logger, setup_logging
 from goose.cli.prompt.goose_prompt_session import GoosePromptSession
 from goose.notifier import Notifier
 from goose.profile import Profile
 from goose.utils import droid, load_plugins
+from goose.utils._cost_calculator import get_total_cost_message
 from goose.utils.session_file import read_from_file, write_to_file
 
 RESUME_MESSAGE = "I see we were interrupted. How can I help you?"
@@ -97,6 +94,7 @@ def __init__(
         self.notifier = SessionNotifier(self.status_indicator)
 
         self.exchange = build_exchange(profile=load_profile(profile), notifier=self.notifier)
+        setup_logging(log_file_directory=LOG_PATH)
 
         if name is not None and self.session_file_path.exists():
             messages = self.load_session()
@@ -173,6 +171,7 @@ def run(self) -> None:
             message = Message.user(text=user_input.text) if user_input.to_continue() else None
 
         self.save_session()
+        get_logger().info(get_total_cost_message(self.exchange.get_token_usage()))
 
     def reply(self) -> None:
         """Reply to the last user message, calling tools as needed

src/goose/utils/_cost_calculator.py

@@ -0,0 +1,39 @@
+from typing import List, Optional
+from exchange.token_usage_collector import TokenUsage
+
+PRICES = {
+    "gpt-4o": (5.00, 15.00),
+    "gpt-4o-2024-08-06": (2.50, 10.00),
+    "gpt-4o-mini": (0.150, 0.600),
+    "gpt-4o-mini-2024-07-18": (0.150, 0.600),
+    "o1-preview": (15.00, 60.00),
+    "o1-mini": (3.00, 12.00),
+    "claude-3-5-sonnet-20240620": (3.00, 15.00),
+    "anthropic.claude-3-5-sonnet-20240620-v1:0": (3.00, 15.00),
+    "claude-3-opus-20240229": (15.00, 75.00),
+    "anthropic.claude-3-opus-20240229-v1:0": (15.00, 75.00),
+    "claude-3-haiku-20240307": (0.25, 1.25),
+    "anthropic.claude-3-haiku-20240307-v1:0": (0.25, 1.25),
+}
+
+
+def _calculate_cost(token_usage: TokenUsage) -> Optional[float]:
+    model_name = token_usage.model.lower()
+    if model_name in PRICES:
+        input_token_price, output_token_price = PRICES[model_name]
+        return (input_token_price * token_usage.input_tokens + output_token_price * token_usage.output_tokens) / 1000000
+    return None
+
+
+def get_total_cost_message(token_usage: List[TokenUsage]) -> str:
+    total_cost = 0
+    message = ""
+    for usage in token_usage:
+        cost = _calculate_cost(usage)
+        if cost is not None:
+            message += f"Cost for {str(usage)}: ${cost:.2f}\n"
+            total_cost += cost
+        else:
+            message += f"Cost for {str(usage)}: Not available\n"
+    message += f"Total cost: ${total_cost:.2f}"
+    return message

tests/utils/test_cost_calculator.py

@@ -0,0 +1,22 @@
+from goose.utils._cost_calculator import _calculate_cost, get_total_cost_message
+from exchange.token_usage_collector import TokenUsage
+
+
+def test_calculate_cost():
+    cost = _calculate_cost(TokenUsage(model="gpt-4o", input_tokens=10000, output_tokens=600))
+    assert cost == 0.059
+
+
+def test_get_total_cost_message():
+    message = get_total_cost_message(
+        [
+            TokenUsage(model="gpt-4o", input_tokens=10000, output_tokens=600),
+            TokenUsage(model="gpt-4o-mini", input_tokens=3000000, output_tokens=4000000),
+        ]
+    )
+    expected_message = (
+        "Cost for TokenUsage(model='gpt-4o', input_tokens=10000, output_tokens=600): $0.06\n"
+        + "Cost for TokenUsage(model='gpt-4o-mini', input_tokens=3000000, output_tokens=4000000)"
+        + ": $2.85\nTotal cost: $2.91"
+    )
+    assert message == expected_message