zach
/
RA.Aid
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(output.py): add get_cost_subtitle function to generate cost subtitles for agent output 

refactor(output.py): simplify subtitle generation in print_agent_output function by using get_cost_subtitle
This commit is contained in:
Ariel Frischer 2025-03-10 16:58:39 -07:00
parent 416689b030
commit f43c5e72b6
1 changed files with 13 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
ra_aid/console/output.py
View File

@ -11,6 +11,13 @@ from ra_aid.callbacks.anthropic_callback_handler import AnthropicCallbackHandler
from .formatting import console 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:
return f"Total Cost: ${cost_cb.total_cost:.6f} | Tokens: {cost_cb.total_tokens}"
return None
def print_agent_output( def print_agent_output(
chunk: Dict[str, Any], chunk: Dict[str, Any],
agent_type: Literal["CiaynAgent", "React"], agent_type: Literal["CiaynAgent", "React"],
@ -30,9 +37,7 @@ def print_agent_output(
if isinstance(msg.content, list): if isinstance(msg.content, list):
for content in msg.content: for content in msg.content:
if content["type"] == "text" and content["text"].strip(): if content["type"] == "text" and content["text"].strip():
subtitle = None subtitle = get_cost_subtitle(cost_cb)
if cost_cb:
subtitle = f"Cost: ${cost_cb.total_cost:.6f} | Tokens: {cost_cb.total_tokens}"
console.print( console.print(
Panel( Panel(
@ -44,9 +49,7 @@ def print_agent_output(
) )
else: else:
if msg.content.strip(): if msg.content.strip():
subtitle = None subtitle = get_cost_subtitle(cost_cb)
if cost_cb:
subtitle = f"Total Cost: ${cost_cb.total_cost:.6f} | Tokens: {cost_cb.total_tokens}"
console.print( console.print(
Panel( Panel(
@ -60,10 +63,14 @@ def print_agent_output(
for msg in chunk["tools"]["messages"]: for msg in chunk["tools"]["messages"]:
if msg.status == "error" and msg.content: if msg.status == "error" and msg.content:
err_msg = msg.content.strip() err_msg = msg.content.strip()
subtitle = get_cost_subtitle(cost_cb)
console.print( console.print(
Panel( Panel(
Markdown(err_msg), Markdown(err_msg),
title="❌ Tool Error", title="❌ Tool Error",
subtitle=subtitle,
subtitle_align="right",
border_style="red bold", border_style="red bold",
) )
) )