From 276e86a0332943f65460b391e8cf071446376528 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 26 Jan 2026 14:20:41 +0200 Subject: [PATCH] Fix citation selection when we call visualize --- haiku_rag_slim/haiku/rag/chat/app.py | 10 ++++------ haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py | 7 +++---- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 463be357..7dc40b79 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -317,8 +317,8 @@ class ChatApp(App): widget.remove_class("selected") def on_descendant_focus(self, _event: object) -> None: - """Clear citation selection when input is focused.""" - if isinstance(self.focused, Input): + """Clear citation selection when chat input is focused.""" + if isinstance(self.focused, Input) and self.focused.id == "chat-input": self._clear_citation_selection() async def action_show_visual(self) -> None: @@ -364,10 +364,8 @@ class ChatApp(App): for widget in chat_history.query(CitationWidget): widget.remove_class("selected") - # Add selected class to the newly selected citation - citation_widgets = list(chat_history.query(CitationWidget)) - if 0 <= event.citation_index < len(citation_widgets): - citation_widgets[event.citation_index].add_class("selected") + # Add selected class to the widget that was focused + event.widget.add_class("selected") async def action_show_filter(self) -> None: """Show document filter modal.""" diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py b/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py index 006a8cfc..328e1087 100644 --- a/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py +++ b/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py @@ -90,9 +90,9 @@ class CitationWidget(Collapsible): class Selected(Message): """Message sent when a citation is selected.""" - def __init__(self, citation_index: int) -> None: + def __init__(self, widget: "CitationWidget") -> None: super().__init__() - self.citation_index = citation_index + self.widget = widget def __init__(self, citation: Citation, **kwargs) -> None: title = f"[{citation.index}] {citation.document_title or citation.document_uri}" @@ -120,8 +120,7 @@ class CitationWidget(Collapsible): def on_focus(self) -> None: """When focused, mark as selected.""" - index = self.citation.index or 1 - self.post_message(self.Selected(index - 1)) + self.post_message(self.Selected(self)) def on_key(self, event: "Key") -> None: """Handle Enter to toggle expand/collapse."""