diff --git a/gui.py b/gui.py index abf9244..dd2dd93 100644 --- a/gui.py +++ b/gui.py @@ -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() diff --git a/tests/test_gui_updates.py b/tests/test_gui_updates.py new file mode 100644 index 0000000..fa658d7 --- /dev/null +++ b/tests/test_gui_updates.py @@ -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")