From ee835b77c2cb3cf3523462576ba21d62c742342a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 28 Aug 2026 15:12:55 +0300 Subject: [PATCH] Select a filtered document from the database that holds it The chat filter modal keys a selection by (database, id): copies of a database share document ids, and checking one copy left the other reading as selected. Applying the filter narrows the question's `sources` to the databases the selection names. A twin id inside another selected database still matches there: a serialized id filter cannot carry a source. --- CHANGELOG.md | 2 +- docs/configuration/storage.md | 4 +- haiku_rag_slim/haiku/rag/chat/app.py | 15 ++- .../rag/chat/widgets/document_filter_modal.py | 58 ++++++---- tests/chat/test_chat_app.py | 107 ++++++++++++++++-- 5 files changed, 147 insertions(+), 39 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ba8fd32d..816cf1f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,7 +24,7 @@ - `haiku-rag settings` prints YAML. - The chat document filter pages results and lists the selected separately. - Selection is by document ID, and a typed search applies on enter. + Selection is by document ID and database, and a typed search applies on enter. - `haiku-rag list` prints only the fields a document has. - `haiku-rag` and `haiku-ingester` exit with a message on an embedder mismatch. - Capabilities created without a client honor `lancedb.uri`. diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index f8dbdf4a..950fec71 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -282,9 +282,7 @@ requires one whenever the client covers a set. The analysis sandbox rejects shared document IDs because its mount path is `/documents/{id}/`. -The chat document filter selects by document ID and applies `id IN (...)` to -every covered database, so selecting an ID that copies share matches the -document in each of them. +The chat document filter selects by document and database: the search is narrowed to the databases the selection names, and the ID filter applies within them. An ID that copies share still matches in every selected database that holds it. #### Ranking diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 442dd3bb..dca3a3da 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -105,7 +105,7 @@ class ChatApp(App): self._state: dict[str, Any] = {} self._is_processing = False self._current_worker: Worker[None] | None = None - self._document_filter: list[str] = [] + self._document_filter: list[tuple[str | None, str]] = [] self._images: list[bytes] = [] # Stable per-launch id for multi-turn model and telemetry correlation. self._conversation_id = str(uuid.uuid4()) @@ -425,12 +425,20 @@ class ChatApp(App): ) def on_document_filter_modal_filter_changed(self, event: Any) -> None: - """Handle document filter changes from modal.""" + """Scope the conversation to the selection: the filter carries the ids, + and `sources` restricts the search to the databases the selection names. + """ from haiku.rag.tools.filters import build_document_id_filter self._document_filter = event.selected - doc_filter = build_document_id_filter(self._document_filter) + doc_filter = build_document_id_filter( + sorted({doc_id for _, doc_id in event.selected}) + ) + selected_sources = {source for source, _ in event.selected} + sources: list[str] | None = None + if selected_sources and None not in selected_sources: + sources = sorted(s for s in selected_sources if s is not None) for namespace, state_type in ( (RAG_STATE_NAMESPACE, RAGState), (ANALYSIS_STATE_NAMESPACE, AnalysisState), @@ -438,4 +446,5 @@ class ChatApp(App): if namespace in self._state: state = state_type.model_validate(self._state[namespace]) state.document_filter = doc_filter + state.sources = sources self._state[namespace] = state.model_dump(mode="json") diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py b/haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py index d06cc1dd..aa4d0e90 100644 --- a/haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py +++ b/haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py @@ -16,26 +16,31 @@ DOCUMENT_PAGE = 200 class DocumentCheckbox(Checkbox): - def __init__(self, label: str, doc_id: str, *, value: bool) -> None: + def __init__( + self, label: str, source: str | None, doc_id: str, *, value: bool + ) -> None: super().__init__(label, value=value, classes="doc-checkbox") + self.source = source self.doc_id = doc_id -def _labelled(docs) -> list[tuple[str, str]]: - """Each document's label and id, sorted. The database is named alongside the - title, which a title alone does not say. Labels are escaped: titles and - database names are data, not Textual markup.""" - return sorted( +def _labelled(docs) -> list[tuple[str, str | None, str]]: + """Each document's label, database and id, sorted by label. The database is + named alongside the title, which a title alone does not say. Labels are + escaped: titles and database names are data, not Textual markup.""" + rows = [ ( escape( f"{doc.title or doc.uri or doc.id}" + (f" ({doc.source})" if doc.source else "") ), + doc.source, doc.id, ) for doc in docs if doc.id is not None - ) + ] + return sorted(rows, key=lambda row: (row[0], row[1] or "", row[2])) def search_filter(term: str) -> str | None: @@ -123,21 +128,22 @@ class DocumentFilterModal(ModalScreen): """ class FilterChanged(Message): - """Emitted when the document filter selection changes.""" + """Emitted when the document filter selection changes. Each selection + names its database, since copies of a database share document ids.""" - def __init__(self, selected: list[str]) -> None: + def __init__(self, selected: list[tuple[str | None, str]]) -> None: super().__init__() self.selected = selected def __init__( self, client: HaikuRAG, - selected: list[str] | None = None, + selected: list[tuple[str | None, str]] | None = None, ) -> None: super().__init__() self.client = client self.initial_selected = selected or [] - self._selected: set[str] = set(self.initial_selected) + self._selected: set[tuple[str | None, str]] = set(self.initial_selected) self._matching = 0 self._search = "" self._page = 0 @@ -180,15 +186,20 @@ class DocumentFilterModal(ModalScreen): self._listing_selected = False if self._listing_selected: - ids = sorted(self._selected) - self._matching = len(ids) - page = ids[self._page * DOCUMENT_PAGE : (self._page + 1) * DOCUMENT_PAGE] + keys = sorted(self._selected, key=lambda key: (key[0] or "", key[1])) + self._matching = len(keys) + page = keys[self._page * DOCUMENT_PAGE : (self._page + 1) * DOCUMENT_PAGE] + page_keys = set(page) docs = ( - list( - await self.client.list_documents( - filter=build_document_id_filter(page) + [ + doc + for doc in await self.client.list_documents( + filter=build_document_id_filter( + sorted({doc_id for _, doc_id in page}) + ) ) - ) + if (doc.source, doc.id) in page_keys + ] if page else [] ) @@ -207,8 +218,10 @@ class DocumentFilterModal(ModalScreen): await filter_list.remove_children() boxes = [ - DocumentCheckbox(label, doc_id, value=doc_id in self._selected) - for label, doc_id in _labelled(docs) + DocumentCheckbox( + label, source, doc_id, value=(source, doc_id) in self._selected + ) + for label, source, doc_id in _labelled(docs) ] if boxes: await filter_list.mount_all(boxes) @@ -253,10 +266,11 @@ class DocumentFilterModal(ModalScreen): if not isinstance(checkbox, DocumentCheckbox): return + key = (checkbox.source, checkbox.doc_id) if event.value: - self._selected.add(checkbox.doc_id) + self._selected.add(key) else: - self._selected.discard(checkbox.doc_id) + self._selected.discard(key) if self._listing_selected: # This listing is the selection, so removing one changes both what diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 4a5caed9..224b08f7 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -463,8 +463,8 @@ async def test_document_filter_updates_rag_state(temp_db_path: Path): async with app.run_test(): # The selection is document ids, so a repeated title cannot widen it. selected = [ - "6f1c2d4e-0000-4000-8000-000000000001", - "6f1c2d4e-0000-4000-8000-000000000002", + (None, "6f1c2d4e-0000-4000-8000-000000000001"), + (None, "6f1c2d4e-0000-4000-8000-000000000002"), ] app.on_document_filter_modal_filter_changed( DocumentFilterModal.FilterChanged(selected) @@ -472,15 +472,57 @@ async def test_document_filter_updates_rag_state(temp_db_path: Path): # RAGState.document_filter should be set rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE]) - expected_filter = build_document_id_filter(selected) + expected_filter = build_document_id_filter( + [doc_id for _, doc_id in selected] + ) assert rag_state.document_filter == expected_filter assert rag_state.document_filter is not None assert "LIKE" not in rag_state.document_filter + # An unnamed database leaves the question unscoped by source. + assert rag_state.sources is None # The state snapshot should also reflect the change assert app._state["rag"]["document_filter"] == expected_filter +@pytest.mark.asyncio +async def test_document_filter_narrows_sources_to_the_selection(temp_db_path: Path): + """The filter carries ids, and `sources` restricts the question to the + databases the selection names.""" + from haiku.rag.chat.app import RAG_STATE_NAMESPACE + from haiku.rag.chat.widgets.document_filter_modal import DocumentFilterModal + + app, mock_client = _make_app_with_state(temp_db_path) + + with ( + patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag, + _covering_returns(_stub_rag, mock_client), + ): + async with app.run_test(): + app.on_document_filter_modal_filter_changed( + DocumentFilterModal.FilterChanged( + [("alpha", "id-one"), ("alpha", "id-two")] + ) + ) + rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE]) + assert rag_state.sources == ["alpha"] + + app.on_document_filter_modal_filter_changed( + DocumentFilterModal.FilterChanged( + [("alpha", "id-one"), ("beta", "id-three")] + ) + ) + rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE]) + assert rag_state.sources == ["alpha", "beta"] + + app.on_document_filter_modal_filter_changed( + DocumentFilterModal.FilterChanged([]) + ) + rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE]) + assert rag_state.sources is None + assert rag_state.document_filter is None + + @pytest.mark.asyncio async def test_document_filter_cleared_when_empty(temp_db_path: Path): """Test that clearing all document filters sets document_filter to None.""" @@ -496,7 +538,7 @@ async def test_document_filter_cleared_when_empty(temp_db_path: Path): async with app.run_test(): # First set a filter app.on_document_filter_modal_filter_changed( - DocumentFilterModal.FilterChanged(["AI Overview"]) + DocumentFilterModal.FilterChanged([(None, "AI Overview")]) ) rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE]) assert rag_state.document_filter is not None @@ -684,7 +726,52 @@ class TestDocumentSelectionIdentity: boxes[0].value = True await pilot.pause() - assert modal._selected == {"id-one"} + assert modal._selected == {("arxiv", "id-one")} + + @pytest.mark.asyncio + async def test_a_shared_id_selects_only_the_named_database_copy( + self, temp_db_path: Path + ): + """Copies of a database share document ids, so a selection carries the + database name and checking one copy leaves the other unselected.""" + from haiku.rag.chat.widgets.document_filter_modal import ( + DocumentCheckbox, + DocumentFilterModal, + ) + from haiku.rag.store.models.document import Document + + client = AsyncMock() + client.covers_multiple = True + client.source_names = ("alpha", "beta") + client.list_documents.return_value = [ + Document(id="id-x", content="", title="Report", source="alpha"), + Document(id="id-x", content="", title="Report", source="beta"), + ] + client.count_documents.return_value = 2 + + modal = DocumentFilterModal(client=client) + app, _ = _make_app(temp_db_path, client) + with ( + patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag, + _covering_returns(_stub_rag, client), + ): + async with app.run_test() as pilot: + await app.push_screen(modal) + await pilot.pause() + + boxes = list(modal.query(DocumentCheckbox)) + assert [str(b.label) for b in boxes] == [ + "Report (alpha)", + "Report (beta)", + ] + + boxes[0].value = True + await pilot.pause() + assert modal._selected == {("alpha", "id-x")} + + await modal._load_documents() + rebuilt = list(modal.query(DocumentCheckbox)) + assert [b.value for b in rebuilt] == [True, False] def test_a_label_that_looks_like_markup_is_text(self): from haiku.rag.chat.widgets.document_filter_modal import ( @@ -699,8 +786,8 @@ class TestDocumentSelectionIdentity: ) ] - ((label, doc_id),) = _labelled(docs) - box = DocumentCheckbox(label, doc_id, value=False) + ((label, source, doc_id),) = _labelled(docs) + box = DocumentCheckbox(label, source, doc_id, value=False) assert str(box.label) == "Report [/red] (alpha [/x])" @@ -952,7 +1039,7 @@ class TestKeepingSelectionsReachable: client.list_documents.side_effect = listing modal = DocumentFilterModal( - client=client, selected=[d.id or "" for d in picked] + client=client, selected=[(None, d.id or "") for d in picked] ) app, _ = _make_app(temp_db_path, client) with ( @@ -1018,7 +1105,7 @@ class TestKeepingSelectionsReachable: client.list_documents.side_effect = listing modal = DocumentFilterModal( - client=client, selected=[d.id or "" for d in picked] + client=client, selected=[(None, d.id or "") for d in picked] ) app, _ = _make_app(temp_db_path, client) with ( @@ -1050,7 +1137,7 @@ class TestKeepingSelectionsReachable: # The row is gone from the listing, not merely unchecked. assert "sel-0200" not in remaining assert len(remaining) == DOCUMENT_PAGE - assert modal._selected == {d.id for d in picked} - {"sel-0200"} + assert modal._selected == {(None, d.id) for d in picked} - {(None, "sel-0200")} # The page it was on no longer exists, so the modal does not report it. assert modal._page == 0 assert "page" not in footer