feat(agent_utils.py): pass callback to print_agent_output to track costs and tokens

feat(output.py): enhance print_agent_output to display cost and token information in the output panel

ra_aid/agent_utils.py

@@ -1591,7 +1591,7 @@ def _run_agent_stream(agent: RAgents, msg_list: list[BaseMessage]):
             logger.debug("Agent output: %s", chunk)
             check_interrupt()
             agent_type = get_agent_type(agent)
-            print_agent_output(chunk, agent_type)
+            print_agent_output(chunk, agent_type, cb)
 
             if is_completed() or should_exit():
                 reset_completion_flags()

ra_aid/console/output.py

@@ -5,13 +5,16 @@ from rich.markdown import Markdown
 from rich.panel import Panel
 
 from ra_aid.exceptions import ToolExecutionError
+from ra_aid.callbacks.anthropic_callback_handler import AnthropicCallbackHandler
 
 # Import shared console instance
 from .formatting import console
 
 
 def print_agent_output(
-    chunk: Dict[str, Any], agent_type: Literal["CiaynAgent", "React"]
+    chunk: Dict[str, Any],
+    agent_type: Literal["CiaynAgent", "React"],
+    cost_cb: Optional[AnthropicCallbackHandler] = None,
 ) -> None:
     """Print only the agent's message content, not tool calls.
 
@@ -27,13 +30,31 @@ def print_agent_output(
                 if isinstance(msg.content, list):
                     for content in msg.content:
                         if content["type"] == "text" and content["text"].strip():
+                            subtitle = None
+                            if cost_cb:
+                                subtitle = f"Cost: ${cost_cb.total_cost:.6f} | Tokens: {cost_cb.total_tokens}"
+
                             console.print(
-                                Panel(Markdown(content["text"]), title="🤖 Assistant")
+                                Panel(
+                                    Markdown(content["text"]),
+                                    title="🤖 Assistant",
+                                    subtitle=subtitle,
+                                    subtitle_align="right",
+                                )
                             )
                 else:
                     if msg.content.strip():
+                        subtitle = None
+                        if cost_cb:
+                            subtitle = f"Total Cost: ${cost_cb.total_cost:.6f} | Tokens: {cost_cb.total_tokens}"
+
                         console.print(
-                            Panel(Markdown(msg.content.strip()), title="🤖 Assistant")
+                            Panel(
+                                Markdown(msg.content.strip()),
+                                title="🤖 Assistant",
+                                subtitle=subtitle,
+                                subtitle_align="right",
+                            )
                         )
     elif "tools" in chunk and "messages" in chunk["tools"]:
         for msg in chunk["tools"]["messages"]: