feat(gui): Implement token budget visualizer

This change adds a progress bar and label to the Provider panel to display the current history token usage against the provider's limit. The UI is updated in real-time.
This commit is contained in:
2026-02-23 13:40:04 -05:00
parent a52f3a2ef8
commit c35170786b
2 changed files with 90 additions and 0 deletions

17
gui.py
View File

@@ -769,6 +769,17 @@ class App:
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 _update_telemetry_panel(self):
"""Updates the token budget visualizer in the Provider panel."""
stats = ai_client.get_history_bleed_stats()
if dpg.does_item_exist("token_budget_bar"):
percentage = stats.get("percentage", 0.0)
dpg.set_value("token_budget_bar", percentage / 100.0 if percentage else 0.0)
if dpg.does_item_exist("token_budget_label"):
current = stats.get("current", 0)
limit = stats.get("limit", 0)
dpg.set_value("token_budget_label", f"{current:,} / {limit:,}")
def _append_comms_entry(self, entry: dict, idx: int):
if not dpg.does_item_exist("comms_scroll"):
return
@@ -1870,6 +1881,11 @@ class App:
callback=self.cb_model_changed,
)
dpg.add_separator()
dpg.add_text("Telemetry")
dpg.add_text("History Token Budget:", color=_LABEL_COLOR)
dpg.add_progress_bar(tag="token_budget_bar", default_value=0.0, width=-1)
dpg.add_text("0 / 0", tag="token_budget_label")
dpg.add_separator()
dpg.add_text("Parameters")
dpg.add_input_float(tag="ai_temperature", label="Temperature", default_value=self.temperature, min_value=0.0, max_value=2.0)
dpg.add_input_int(tag="ai_max_tokens", label="Max Tokens (Output)", default_value=self.max_tokens, step=1024)
@@ -2207,6 +2223,7 @@ class App:
# Flush any comms entries queued from background threads
self._flush_pending_comms()
self._update_telemetry_panel()
dpg.render_dearpygui_frame()

73
tests/test_gui_updates.py Normal file
View File

@@ -0,0 +1,73 @@
import pytest
from unittest.mock import patch, MagicMock
import importlib.util
import sys
import dearpygui.dearpygui as dpg
# Load gui.py as a module for testing
spec = importlib.util.spec_from_file_location("gui", "gui.py")
gui = importlib.util.module_from_spec(spec)
sys.modules["gui"] = gui
spec.loader.exec_module(gui)
from gui import App
@pytest.fixture
def app_instance():
"""
Fixture to create an instance of the App class for testing.
It creates a real DPG context but mocks functions that would
render a window or block execution.
"""
dpg.create_context()
# Patch only the functions that would show a window or block,
# and the App methods that rebuild UI on init.
with patch('dearpygui.dearpygui.create_viewport'), \
patch('dearpygui.dearpygui.setup_dearpygui'), \
patch('dearpygui.dearpygui.show_viewport'), \
patch('dearpygui.dearpygui.start_dearpygui'), \
patch('gui.load_config', return_value={}), \
patch.object(App, '_rebuild_files_list'), \
patch.object(App, '_rebuild_shots_list'), \
patch.object(App, '_rebuild_disc_list'), \
patch.object(App, '_rebuild_disc_roles_list'), \
patch.object(App, '_rebuild_discussion_selector'), \
patch.object(App, '_refresh_project_widgets'):
app = App()
yield app
dpg.destroy_context()
def test_telemetry_panel_updates_correctly(app_instance):
"""
Tests that the _update_telemetry_panel method correctly updates
DPG widgets based on the stats from ai_client.
"""
# 1. Define the mock stats to be returned by ai_client
mock_stats = {
"provider": "anthropic",
"limit": 180000,
"current": 135000,
"percentage": 75.0,
}
# 2. Patch the dependencies
with patch('ai_client.get_history_bleed_stats', return_value=mock_stats) as mock_get_stats, \
patch('dearpygui.dearpygui.set_value') as mock_set_value, \
patch('dearpygui.dearpygui.does_item_exist', return_value=True) as mock_does_item_exist:
# 3. Call the method under test
app_instance._update_telemetry_panel()
# 4. Assert the results
mock_get_stats.assert_called_once()
# Check that the code verified the existence of the widgets
mock_does_item_exist.assert_any_call("token_budget_bar")
mock_does_item_exist.assert_any_call("token_budget_label")
# Check that the widgets were updated with the correct values
mock_set_value.assert_any_call("token_budget_bar", 0.75)
# Note: My implementation formats numbers with commas
mock_set_value.assert_any_call("token_budget_label", "135,000 / 180,000")