From 12d27952d52e0b5acc6ea59320e6d93f14f3df95 Mon Sep 17 00:00:00 2001
From: AI Christianson <ai.christianson@christianson.ai>
Date: Wed, 12 Mar 2025 10:21:06 -0400
Subject: [PATCH] add --show-cost flag

---
 ra_aid/__main__.py                                | 7 +++++++
 ra_aid/config.py                                  | 3 ++-
 ra_aid/console/output.py                          | 8 ++++++--
 ra_aid/database/repositories/config_repository.py | 2 ++
 4 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py
index aaf6f84..3121e5d 100644
--- a/ra_aid/__main__.py
+++ b/ra_aid/__main__.py
@@ -293,6 +293,11 @@ Examples:
         action="store_true",
         help="Display model thinking content extracted from think tags when supported by the model",
     )
+    parser.add_argument(
+        "--show-cost",
+        action="store_true",
+        help="Display cost information as the agent works",
+    )
     parser.add_argument(
         "--reasoning-assistance",
         action="store_true",
@@ -602,6 +607,7 @@ def main():
                 config_repo.set("experimental_fallback_handler", args.experimental_fallback_handler)
                 config_repo.set("web_research_enabled", web_research_enabled)
                 config_repo.set("show_thoughts", args.show_thoughts)
+                config_repo.set("show_cost", args.show_cost)
                 config_repo.set("force_reasoning_assistance", args.reasoning_assistance)
                 config_repo.set("disable_reasoning_assistance", args.no_reasoning_assistance)
 
@@ -707,6 +713,7 @@ def main():
                     config_repo.set("expert_model", args.expert_model)
                     config_repo.set("temperature", args.temperature)
                     config_repo.set("show_thoughts", args.show_thoughts)
+                    config_repo.set("show_cost", args.show_cost)
                     config_repo.set("force_reasoning_assistance", args.reasoning_assistance)
                     config_repo.set("disable_reasoning_assistance", args.no_reasoning_assistance)
 
diff --git a/ra_aid/config.py b/ra_aid/config.py
index c414f3b..936fc26 100644
--- a/ra_aid/config.py
+++ b/ra_aid/config.py
@@ -6,6 +6,7 @@ DEFAULT_MAX_TOOL_FAILURES = 3
 FALLBACK_TOOL_MODEL_LIMIT = 5
 RETRY_FALLBACK_COUNT = 3
 DEFAULT_TEST_CMD_TIMEOUT = 60 * 5  # 5 minutes in seconds
+DEFAULT_SHOW_COST = False
 
 
 VALID_PROVIDERS = [
@@ -15,4 +16,4 @@ VALID_PROVIDERS = [
     "openai-compatible",
     "deepseek",
     "gemini",
-]
+]
\ No newline at end of file
diff --git a/ra_aid/console/output.py b/ra_aid/console/output.py
index 8a45fec..68b0b85 100644
--- a/ra_aid/console/output.py
+++ b/ra_aid/console/output.py
@@ -6,14 +6,18 @@ from rich.panel import Panel
 
 from ra_aid.exceptions import ToolExecutionError
 from ra_aid.callbacks.anthropic_callback_handler import AnthropicCallbackHandler
+from ra_aid.database.repositories.config_repository import get_config_repository
+from ra_aid.config import DEFAULT_SHOW_COST
 
 # Import shared console instance
 from .formatting import console
 
 
 def get_cost_subtitle(cost_cb: Optional[AnthropicCallbackHandler]) -> Optional[str]:
-    """Generate a subtitle with cost information if a callback is provided."""
-    if cost_cb:
+    """Generate a subtitle with cost information if a callback is provided and show_cost is enabled."""
+    # Only show cost information if both cost_cb is provided AND show_cost is True
+    show_cost = get_config_repository().get("show_cost", DEFAULT_SHOW_COST)
+    if cost_cb and show_cost:
         return f"Cost: ${cost_cb.total_cost:.6f} | Tokens: {cost_cb.total_tokens}"
     return None
 
diff --git a/ra_aid/database/repositories/config_repository.py b/ra_aid/database/repositories/config_repository.py
index 70be4cc..7aa05f3 100644
--- a/ra_aid/database/repositories/config_repository.py
+++ b/ra_aid/database/repositories/config_repository.py
@@ -32,6 +32,7 @@ class ConfigRepository:
             FALLBACK_TOOL_MODEL_LIMIT,
             RETRY_FALLBACK_COUNT,
             DEFAULT_TEST_CMD_TIMEOUT,
+            DEFAULT_SHOW_COST,
             VALID_PROVIDERS,
         )
         
@@ -42,6 +43,7 @@ class ConfigRepository:
             "fallback_tool_model_limit": FALLBACK_TOOL_MODEL_LIMIT,
             "retry_fallback_count": RETRY_FALLBACK_COUNT,
             "test_cmd_timeout": DEFAULT_TEST_CMD_TIMEOUT,
+            "show_cost": DEFAULT_SHOW_COST,
             "valid_providers": VALID_PROVIDERS,
         }