diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 0d110435..5465c532 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -317,7 +317,11 @@ class ChatApp(App): blobs ) - await chat_history.add_citations(citations, picture_bytes=picture_bytes) + await chat_history.add_citations( + citations, + picture_bytes=picture_bytes, + include_collection=self.client is not None and self.client.covers_multiple, + ) if analysis_data := self._state.get(ANALYSIS_STATE_NAMESPACE): analysis_state = AnalysisState.model_validate(analysis_data) 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 da42d634..27cde368 100644 --- a/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py +++ b/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py @@ -118,9 +118,12 @@ class CitationWidget(Collapsible): self, citation: Citation, picture_bytes: list[bytes] | None = None, + include_collection: bool = False, **kwargs, ) -> None: title = f"[{citation.index}] {citation.document_title or citation.document_uri}" + if include_collection and citation.source: + title += f" — {citation.source}" if citation.page_numbers: pages = ", ".join(map(str, citation.page_numbers[:3])) if len(citation.page_numbers) > 3: @@ -432,14 +435,13 @@ class ChatHistory(VerticalScroll): self, citations: list[Citation], picture_bytes: dict[tuple[str | None, str | None], list[bytes]] | None = None, + include_collection: bool = False, ) -> None: """Add citations inline after a response. ``picture_bytes`` maps a citation's ``(source, chunk_id)`` → list of raw - PNG bytes, - one per entry in the citation's ``picture_refs``. Pre-fetched by the - caller (typically the chat app's post-response hook) so widget - construction stays synchronous. + PNG bytes, one per entry in the citation's ``picture_refs``. Pre-fetched + by the caller so widget construction stays synchronous. """ if not citations: return @@ -451,6 +453,7 @@ class ChatHistory(VerticalScroll): picture_bytes=picture_bytes.get( qualified_id(citation.source, citation.chunk_id) ), + include_collection=include_collection, ) await self.mount(widget) self.scroll_end(animate=False) diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 5408195c..1f558dd9 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -806,6 +806,67 @@ class TestRenderingUnattributedPictures: ] +class TestNamingACitationsCollection: + @staticmethod + async def _titles(temp_db_path, covering, *sources: str | None) -> list[str]: + """The collapsed titles of one citation per source, all named alike.""" + from haiku.rag.chat.app import RAG_STATE_NAMESPACE + from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget + from haiku.rag.store.models.citation import Citation + + index = { + f"c{position}": Citation( + document_id=f"d{position}", + chunk_id=f"c{position}", + source=source, + content="body", + document_uri="test://report", + document_title="Quarterly report", + ).model_dump(mode="json") + for position, source in enumerate(sources) + } + + app, _ = _make_app(temp_db_path, covering) + with ( + patch("haiku.rag.chat.app.HaikuRAG") as stub, + _covering_returns(stub, covering), + ): + async with app.run_test() as pilot: + app._state[RAG_STATE_NAMESPACE] = { + "citations": list(index), + "citation_index": index, + } + await app._show_citations_and_programs(app.query_one(ChatHistory)) + await pilot.pause() + return [str(w.title) for w in app.query(CitationWidget)] + + @pytest.mark.asyncio + async def test_one_title_in_two_collections_reads_as_two_citations( + self, temp_db_path: Path + ): + covering = _make_mock_client() + covering.covers_multiple = True + covering.source_names = ("alpha", "beta") + + titles = await self._titles(temp_db_path, covering, "alpha", "beta") + + assert len(set(titles)) == 2 + assert "alpha" in titles[0] + assert "beta" in titles[1] + + @pytest.mark.asyncio + async def test_one_named_database_is_not_named_on_its_citations( + self, temp_db_path: Path + ): + covering = _make_mock_client() + covering.covers_multiple = False + covering.source_names = ("alpha",) + + [title] = await self._titles(temp_db_path, covering, "alpha") + + assert "alpha" not in title + + class TestKeepingSelectionsReachable: """A selection applies whether or not the page shows it, and a checkbox is the only way to remove one."""