danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Select a filtered document from the database that holds it

Browse source 
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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-28 15:12:55 +03:00
parent b601489896
commit ee835b77c2
No known key found for this signature in database
5 changed files with 147 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
CHANGELOG.md
View file

@ -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`.

4
docs/configuration/storage.md
View file

@ -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

15
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -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")

58
haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py
View file

@ -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

107
tests/chat/test_chat_app.py
View file

@ -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