log token usage after every single step

Author: raycai420

aide/agent.py

@@ -3,6 +3,7 @@
 import random
 import time
 from typing import Any, Callable, cast
+import json
 
 import humanize
 from .backend import FunctionSpec, query, TokenCounter
@@ -360,6 +361,19 @@ def step(self, exec_callback: ExecCallbackType) -> bool:
                 logger.info(f"Node {best_node.id} is still the best node")
         self.current_step += 1
 
+        total_cost = self.token_counter.cost()
+        input_tokens_dict = self.token_counter.total_input_tokens
+        output_tokens_dict = self.token_counter.total_output_tokens
+        logger.info(
+            "Total input tokens after current step: %s",
+            json.dumps(input_tokens_dict, sort_keys=True),
+        )
+        logger.info(
+            "Total output tokens after current step: %s",
+            json.dumps(output_tokens_dict, sort_keys=True),
+        )
+        logger.info(f"Total cost after current step: {total_cost:.2f} USD")
+
         # check early exit required by token counter
         if self.token_counter.cost_limit:
             exceed_budget_limit = self.token_counter.exceed_budget_limit()

aide/run.py

@@ -2,7 +2,6 @@
 import logging
 import shutil
 import sys
-import json
 
 from . import backend
 
@@ -198,18 +197,6 @@ def generate_live():
     if exceed_budget_limit and global_step < cfg.agent.steps:
         logger.info(journal_to_string_tree(journal))
         logger.info("Budget limit exceeded, stopping run.")
-    total_cost = agent.token_counter.cost()
-    input_tokens_dict = agent.token_counter.total_input_tokens
-    output_tokens_dict = agent.token_counter.total_output_tokens
-    logger.info(
-        "Total input tokens before report generation: %s",
-        json.dumps(input_tokens_dict, sort_keys=True),
-    )
-    logger.info(
-        "Total output tokens before report generation: %s",
-        json.dumps(output_tokens_dict, sort_keys=True),
-    )
-    logger.info(f"Total cost before report generation: {total_cost:.2f} USD")
     interpreter.cleanup_session()
 
     if cfg.generate_report: