Fix visual grounding citation selection in TUI

This commit is contained in:
Yiorgis Gozadinos 2026-01-26 10:58:33 +02:00
parent ff92ee8271
commit 729a3910d8
No known key found for this signature in database

View file

@ -103,7 +103,6 @@ class ChatApp(App):
self._is_processing = False self._is_processing = False
self._tool_call_widgets: dict[str, Any] = {} self._tool_call_widgets: dict[str, Any] = {}
self._last_citations: list[Citation] = [] self._last_citations: list[Citation] = []
self._selected_citation_idx: int | None = None
self._current_worker: Worker[None] | None = None self._current_worker: Worker[None] | None = None
self._message_history: list[ModelMessage] = [] self._message_history: list[ModelMessage] = []
@ -200,7 +199,6 @@ class ChatApp(App):
# Clear for new query # Clear for new query
self._tool_call_widgets.clear() self._tool_call_widgets.clear()
self._last_citations.clear() self._last_citations.clear()
self._selected_citation_idx = None
# Run agent in a worker to keep UI responsive # Run agent in a worker to keep UI responsive
self._is_processing = True self._is_processing = True
@ -271,7 +269,6 @@ class ChatApp(App):
chat_history = self.query_one(ChatHistory) chat_history = self.query_one(ChatHistory)
await chat_history.clear_messages() await chat_history.clear_messages()
self._last_citations.clear() self._last_citations.clear()
self._selected_citation_idx = None
self._message_history.clear() self._message_history.clear()
# Reset session state for fresh conversation # Reset session state for fresh conversation
self.session_state = ChatSessionState( self.session_state = ChatSessionState(
@ -289,7 +286,6 @@ class ChatApp(App):
chat_history = self.query_one(ChatHistory) chat_history = self.query_one(ChatHistory)
for widget in chat_history.query(CitationWidget): for widget in chat_history.query(CitationWidget):
widget.remove_class("selected") widget.remove_class("selected")
self._selected_citation_idx = None
def on_descendant_focus(self, _event: object) -> None: def on_descendant_focus(self, _event: object) -> None:
"""Clear citation selection when input is focused.""" """Clear citation selection when input is focused."""
@ -298,15 +294,16 @@ class ChatApp(App):
async def action_show_visual(self) -> None: async def action_show_visual(self) -> None:
"""Show visual grounding for the selected citation.""" """Show visual grounding for the selected citation."""
if not self.client or not self._last_citations: if not self.client:
return return
idx = ( # Get citation from selected widget directly
self._selected_citation_idx chat_history = self.query_one(ChatHistory)
if self._selected_citation_idx is not None selected_widgets = list(chat_history.query(CitationWidget).filter(".selected"))
else 0 if not selected_widgets:
) return
citation = self._last_citations[idx]
citation = selected_widgets[0].citation
chunk = await self.client.chunk_repository.get_by_id(citation.chunk_id) chunk = await self.client.chunk_repository.get_by_id(citation.chunk_id)
if not chunk: if not chunk:
return return
@ -342,4 +339,3 @@ class ChatApp(App):
citation_widgets = list(chat_history.query(CitationWidget)) citation_widgets = list(chat_history.query(CitationWidget))
if 0 <= event.citation_index < len(citation_widgets): if 0 <= event.citation_index < len(citation_widgets):
citation_widgets[event.citation_index].add_class("selected") citation_widgets[event.citation_index].add_class("selected")
self._selected_citation_idx = event.citation_index