feat(ui): Implement token usage summary widget

This commit is contained in:
2026-02-23 10:59:29 -05:00
parent 4479c38395
commit e34ff7ef79
3 changed files with 67 additions and 0 deletions

27
gui.py
View File

@@ -47,6 +47,21 @@ def hide_tk_root() -> Tk:
root.wm_attributes("-topmost", True)
return root
def get_total_token_usage() -> dict:
"""Returns aggregated token usage across the entire session from comms log."""
usage = {
"input_tokens": 0,
"output_tokens": 0,
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0
}
for entry in ai_client.get_comms_log():
if entry.get("kind") == "response" and "usage" in entry.get("payload", {}):
u = entry["payload"]["usage"]
for k in usage.keys():
usage[k] += u.get(k, 0) or 0
return usage
# ------------------------------------------------------------------ comms rendering helpers
@@ -713,6 +728,15 @@ class App:
for entry in entries:
self._comms_entry_count += 1
self._append_comms_entry(entry, self._comms_entry_count)
if entries:
self._update_token_usage()
def _update_token_usage(self):
if not dpg.does_item_exist("ai_token_usage"):
return
usage = get_total_token_usage()
total = usage["input_tokens"] + usage["output_tokens"]
dpg.set_value("ai_token_usage", f"Tokens: {total} (In: {usage['input_tokens']} Out: {usage['output_tokens']})")
def _append_comms_entry(self, entry: dict, idx: int):
if not dpg.does_item_exist("comms_scroll"):
@@ -1217,6 +1241,7 @@ class App:
with self._pending_comms_lock:
self._pending_comms.clear()
self._comms_entry_count = 0
self._update_token_usage()
if dpg.does_item_exist("comms_scroll"):
dpg.delete_item("comms_scroll", children_only=True)
@@ -1864,6 +1889,8 @@ class App:
with dpg.group(horizontal=True):
dpg.add_text("Status: idle", tag="ai_status", color=(200, 220, 160))
dpg.add_spacer(width=16)
dpg.add_text("Tokens: 0 (In: 0 Out: 0)", tag="ai_token_usage", color=(180, 255, 180))
dpg.add_spacer(width=16)
dpg.add_button(label="Clear", callback=self.cb_clear_comms)
dpg.add_separator()
with dpg.group(horizontal=True):