diff options
| -rw-r--r-- | llamachat/ui.py | 57 |
1 files changed, 57 insertions, 0 deletions
diff --git a/llamachat/ui.py b/llamachat/ui.py index 4dd758b..189d59a 100644 --- a/llamachat/ui.py +++ b/llamachat/ui.py @@ -459,6 +459,9 @@ class ChatWindow(QMainWindow): self.assistant_buffer = "" self.reasoning_buffer = "" self.searches: list[dict] = [] + self.turn_prompt_tokens: int | None = None + self.turn_completion_tokens: int | None = None + self._conversation_cost = 0.0 # Every rendered bubble, so a reasoning toggle can redraw the # transcript without refetching anything. self.bubbles: list[dict] = [] @@ -539,6 +542,9 @@ class ChatWindow(QMainWindow): self.meter = ContextMeter() top.addWidget(self.meter) + self.cost = CostLabel() + top.addWidget(self.cost) + self.search_box = QLineEdit() self.search_box.setPlaceholderText("Search history… (Ctrl+F)") self.search_box.setClearButtonEnabled(True) @@ -779,6 +785,7 @@ class ChatWindow(QMainWindow): self.show_status("; ".join(all_problems), error=True) else: self.hide_status() + self.update_cost() def current_model(self) -> str: return self.model_box.currentData() or self.model_box.currentText() @@ -941,10 +948,12 @@ class ChatWindow(QMainWindow): self.history_list.clearSelection() self.hide_status() self.select_prompt(self.cfg.default_prompt) + self._conversation_cost = 0.0 self.update_meter() # Starting a conversation implies wanting to type in it, and this # also takes focus back out of the search box. self.input.setFocus() + self.update_cost() # -- attachments ------------------------------------------------------ @@ -1139,6 +1148,7 @@ class ChatWindow(QMainWindow): """ if not hasattr(self, "meter"): return # still building the window + self.update_cost() info = self.current_info() limit = info.ctx_size or 0 @@ -1167,6 +1177,36 @@ class ChatWindow(QMainWindow): used = backend.estimate_tokens(messages, self.cfg.chars_per_token) self.meter.set_usage(used, limit, exact=False) + def update_cost(self) -> None: + """Refresh the money readout from stored counts plus the draft.""" + if not hasattr(self, "cost"): + return # still building the window + model_id = self.current_model() + billable = models_mod.is_billable(model_id, self.cfg.providers) + priced = models_mod.is_priced(model_id, self.cfg.providers, self.store) + if not billable or not priced: + self.cost.set_cost(0.0, 0.0, billable, priced) + return + + # Reopening a conversation resends its whole history, so the + # projection has to price everything that would go out, not just + # what was typed. That is what makes an expensive turn visible + # before it is paid rather than after. + pending = 0 + if self.session_id is not None or self.input.toPlainText(): + pending = backend.estimate_tokens( + self._chat_context( + backend.build_user_content( + self.input.toPlainText(), self.attachments + ) + ), + self.cfg.chars_per_token, + ) + projected = models_mod.projected_cost( + pending, model_id, self.cfg.providers, self.store + ) + self.cost.set_cost(self._conversation_cost, projected, billable, priced) + def _search_config(self) -> SearchConfig: return SearchConfig( enabled=self.cfg.search_enabled, @@ -1218,7 +1258,10 @@ class ChatWindow(QMainWindow): limit = info.ctx_size or 0 # What the next turn starts from is everything sent plus the reply. self.exact_tokens = total_tokens or prompt_tokens + self.turn_prompt_tokens = prompt_tokens + self.turn_completion_tokens = max(total_tokens - prompt_tokens, 0) self.meter.set_usage(self.exact_tokens, limit, exact=True) + self.update_cost() @Slot(str) def _on_reasoning(self, piece: str) -> None: @@ -1254,12 +1297,20 @@ class ChatWindow(QMainWindow): self.assistant_buffer, self.reasoning_buffer, self._searches_json(), + prompt_tokens=self.turn_prompt_tokens, + completion_tokens=self.turn_completion_tokens, + model=self.current_model(), + ) + rows = self.history.messages(self.session_id) + self._conversation_cost = models_mod.conversation_cost( + rows, self.cfg.providers, self.store ) titled = self.assistant_message_id is not None and self._should_title() self._teardown_stream() if titled: self._start_titling() self.refresh_history() + self.update_cost() def _should_title(self) -> bool: session = self.history.get_session(self.session_id) @@ -1329,6 +1380,8 @@ class ChatWindow(QMainWindow): self.thread = None self.worker = None self.assistant_message_id = None + self.turn_prompt_tokens = None + self.turn_completion_tokens = None self.send_button.setEnabled(True) self.stop_button.hide() @@ -1521,6 +1574,9 @@ class ChatWindow(QMainWindow): "searches": _searches(_column(row, "searches")), } ) + self._conversation_cost = models_mod.conversation_cost( + self.history.messages(session_id), self.cfg.providers, self.store + ) self._render() self._scroll_to_end() @@ -1534,6 +1590,7 @@ class ChatWindow(QMainWindow): ) else: self.hide_status() + self.update_cost() # -- status line ------------------------------------------------------ |