d93cb5c891
`reader_for` returned itself before looking at the name, so an "alpha" client
answered `reader_for("beta")` with alpha's reader. A citation naming another
database would have been read from the wrong one.
Both paths decide through `clients_covering` now, so one database refuses a
wrong name the way a set already did.
The chat citation test said mounting lends the client without showing it; a
chat test shows it and the docstring points there.
779 lines
28 KiB
Python
779 lines
28 KiB
Python
from copy import deepcopy
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
from typer.testing import CliRunner
|
|
|
|
from haiku.rag.capabilities.rag import RAGState, create_capability
|
|
from haiku.rag.cli import _cli as cli
|
|
from tests.conftest import _covering_returns, for_path
|
|
|
|
runner = CliRunner()
|
|
|
|
|
|
def test_chat_command():
|
|
"""Test chat command launches chat TUI."""
|
|
with patch("haiku.rag.chat.run_chat") as mock_chat:
|
|
mock_chat.return_value = None
|
|
|
|
result = runner.invoke(cli, ["chat"])
|
|
|
|
assert result.exit_code == 0
|
|
mock_chat.assert_called_once()
|
|
|
|
|
|
def test_run_chat_creates_app_and_runs(temp_db_path: Path):
|
|
"""Test run_chat() eagerly attaches one capability and runs the app."""
|
|
with patch("haiku.rag.chat.app.ChatApp") as mock_app:
|
|
from haiku.rag.chat import run_chat
|
|
|
|
run_chat(db_path=temp_db_path)
|
|
|
|
mock_app.return_value.run.assert_called_once()
|
|
attached = mock_app.call_args.kwargs["capabilities"]
|
|
assert len(attached) == 1
|
|
assert attached[0].defer_loading is False
|
|
|
|
|
|
def test_run_chat_covers_a_configured_set(tmp_path, monkeypatch):
|
|
"""A configured set has no single path to fall back to: the app is handed
|
|
None so the client opens the set."""
|
|
import haiku.rag.config as config_module
|
|
from haiku.rag.config import AppConfig, set_config
|
|
from haiku.rag.config.models import LanceDBConfig
|
|
|
|
monkeypatch.setattr(config_module, "_config", None)
|
|
set_config(
|
|
AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
databases={
|
|
"a": str(tmp_path / "a.lancedb"),
|
|
"b": str(tmp_path / "b.lancedb"),
|
|
}
|
|
)
|
|
)
|
|
)
|
|
with patch("haiku.rag.chat.app.ChatApp") as app:
|
|
from haiku.rag.chat import run_chat
|
|
|
|
run_chat(db_path=None)
|
|
|
|
assert app.call_args.kwargs["scope"].names == ("a", "b")
|
|
|
|
|
|
def test_chat_capabilities_read_the_named_database(tmp_path, monkeypatch):
|
|
"""`--db PATH` and `--db-name NAME` have to reach the capabilities too, or
|
|
they answer from the default database or the whole set."""
|
|
import haiku.rag.config as config_module
|
|
from haiku.rag.client.scope import DatabaseScope
|
|
from haiku.rag.config import set_config
|
|
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
|
|
|
monkeypatch.setattr(config_module, "_config", None)
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
databases={
|
|
"a": str(tmp_path / "a.lancedb"),
|
|
"b": str(tmp_path / "b.lancedb"),
|
|
}
|
|
)
|
|
)
|
|
set_config(config)
|
|
|
|
with patch("haiku.rag.chat.app.ChatApp") as chat_app:
|
|
from haiku.rag.chat import run_chat
|
|
|
|
run_chat(scope=DatabaseScope.resolve(config, database_name="b"))
|
|
[named] = chat_app.call_args.kwargs["capabilities"]
|
|
|
|
run_chat(scope=DatabaseScope.resolve(config))
|
|
[covering] = chat_app.call_args.kwargs["capabilities"]
|
|
|
|
# The chat lends its own client, so this scope is the fallback: what matters
|
|
# is that it places the named database rather than the whole set.
|
|
[placed] = named.scope.databases
|
|
assert placed.db_path == tmp_path / "b.lancedb"
|
|
assert named.config.lancedb.databases == {}
|
|
assert covering.scope.names == ("a", "b")
|
|
assert set(covering.config.lancedb.databases) == {"a", "b"}
|
|
|
|
|
|
def test_run_chat_defers_multiple_capabilities(temp_db_path: Path):
|
|
"""Test chat only defers capabilities when routing between multiple choices."""
|
|
with patch("haiku.rag.chat.app.ChatApp") as mock_app:
|
|
from haiku.rag.chat import run_chat
|
|
|
|
run_chat(db_path=temp_db_path, capabilities=["rag", "analysis"])
|
|
|
|
attached = mock_app.call_args.kwargs["capabilities"]
|
|
assert len(attached) == 2
|
|
assert all(capability.defer_loading for capability in attached)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("enabled", "expected_model", "expected_vision"),
|
|
[
|
|
(["analysis"], "analysis-model", False),
|
|
(["rag"], "qa-model", True),
|
|
(["rag", "analysis"], "qa-model", True),
|
|
],
|
|
)
|
|
def test_run_chat_gates_capability_vision_on_driving_model(
|
|
temp_db_path: Path, enabled, expected_model, expected_vision
|
|
):
|
|
"""Analysis-only chat runs on analysis.model; otherwise on qa.model. Every
|
|
attached capability's vision gate tracks that one driving model."""
|
|
from haiku.rag.config.models import AppConfig, ModelConfig
|
|
|
|
config = AppConfig()
|
|
config.qa.model = ModelConfig(provider="openai", name="qa-model", vision=True)
|
|
config.analysis.model = ModelConfig(
|
|
provider="openai", name="analysis-model", vision=False
|
|
)
|
|
captured: dict[str, str] = {}
|
|
|
|
def fake_get_model(model_config, _config):
|
|
captured["name"] = model_config.name
|
|
return "resolved-model"
|
|
|
|
with (
|
|
patch("haiku.rag.chat.app.ChatApp") as mock_app,
|
|
patch("haiku.rag.config.get_config", return_value=config),
|
|
patch("haiku.rag.utils.get_model", side_effect=fake_get_model),
|
|
):
|
|
from haiku.rag.chat import run_chat
|
|
|
|
run_chat(db_path=temp_db_path, capabilities=enabled)
|
|
|
|
assert captured["name"] == expected_model
|
|
attached = mock_app.call_args.kwargs["capabilities"]
|
|
assert {capability.vision for capability in attached} == {expected_vision}
|
|
|
|
|
|
def _make_mock_client():
|
|
"""Create a mock HaikuRAG client."""
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
|
# Covers one database: a bare AsyncMock answers `covers_multiple` with a truthy
|
|
# Mock, which would send every read down the covering-a-set branch.
|
|
mock_client.covers_multiple = False
|
|
mock_client.source_names = ()
|
|
mock_client.source = None
|
|
return mock_client
|
|
|
|
|
|
def _make_app(db_path: Path, mock_client: AsyncMock | None = None):
|
|
"""Create a ChatApp with mocked HaikuRAG."""
|
|
from haiku.rag.chat.app import ChatApp
|
|
|
|
if mock_client is None:
|
|
mock_client = _make_mock_client()
|
|
|
|
return ChatApp(
|
|
scope=for_path(db_path),
|
|
capabilities=[create_capability(db_path=db_path)],
|
|
read_only=True,
|
|
), mock_client
|
|
|
|
|
|
def _make_app_with_state(db_path: Path, mock_client: AsyncMock | None = None):
|
|
"""Create a ChatApp with a RAG capability and state."""
|
|
from haiku.rag.chat.app import ChatApp
|
|
|
|
if mock_client is None:
|
|
mock_client = _make_mock_client()
|
|
|
|
return ChatApp(
|
|
scope=for_path(db_path),
|
|
capabilities=[create_capability(db_path=db_path)],
|
|
read_only=True,
|
|
), mock_client
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_app_has_required_widgets(temp_db_path: Path):
|
|
"""Test that ChatApp has the required widgets: ChatHistory, FlexibleInput."""
|
|
from haiku.rag.chat.widgets.chat_history import ChatHistory
|
|
|
|
app, mock_client = _make_app(temp_db_path)
|
|
|
|
with (
|
|
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
|
|
_covering_returns(_stub_rag, mock_client),
|
|
):
|
|
async with app.run_test():
|
|
chat_history = app.query_one(ChatHistory)
|
|
assert chat_history is not None
|
|
|
|
from haiku.rag.chat.widgets.prompt import FlexibleInput
|
|
|
|
chat_input = app.query_one(FlexibleInput)
|
|
assert chat_input is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_app_quit_binding(temp_db_path: Path):
|
|
"""Test that pressing ctrl+q quits the app."""
|
|
app, mock_client = _make_app(temp_db_path)
|
|
|
|
with (
|
|
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
|
|
_covering_returns(_stub_rag, mock_client),
|
|
):
|
|
async with app.run_test() as pilot:
|
|
assert app.is_running
|
|
await pilot.press("ctrl+q")
|
|
assert not app.is_running
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_history_can_add_message(temp_db_path: Path):
|
|
"""Test that ChatHistory can display messages."""
|
|
from haiku.rag.chat.widgets.chat_history import ChatHistory
|
|
|
|
app, mock_client = _make_app(temp_db_path)
|
|
|
|
with (
|
|
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
|
|
_covering_returns(_stub_rag, mock_client),
|
|
):
|
|
async with app.run_test():
|
|
chat_history = app.query_one(ChatHistory)
|
|
|
|
await chat_history.add_message("user", "Hello, how are you?")
|
|
assert len(chat_history.messages) == 1
|
|
assert chat_history.messages[0] == ("user", "Hello, how are you?")
|
|
|
|
await chat_history.add_message("assistant", "I'm doing well, thank you!")
|
|
assert len(chat_history.messages) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_history_can_add_tool_calls(temp_db_path: Path):
|
|
"""Test that ChatHistory can display inline tool calls."""
|
|
from haiku.rag.chat.widgets.chat_history import ChatHistory, ToolCallWidget
|
|
|
|
app, mock_client = _make_app(temp_db_path)
|
|
|
|
with (
|
|
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
|
|
_covering_returns(_stub_rag, mock_client),
|
|
):
|
|
async with app.run_test():
|
|
chat_history = app.query_one(ChatHistory)
|
|
|
|
tool_widget = await chat_history.add_tool_call(
|
|
"tool-1", "search", {"query": "test"}
|
|
)
|
|
assert isinstance(tool_widget, ToolCallWidget)
|
|
assert tool_widget._completed is False
|
|
|
|
chat_history.mark_tool_complete("tool-1")
|
|
assert tool_widget._completed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_history_can_add_citations(temp_db_path: Path):
|
|
"""Test that ChatHistory can display inline citations."""
|
|
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
|
|
from haiku.rag.store.models.citation import Citation
|
|
|
|
app, mock_client = _make_app(temp_db_path)
|
|
|
|
with (
|
|
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
|
|
_covering_returns(_stub_rag, mock_client),
|
|
):
|
|
async with app.run_test():
|
|
chat_history = app.query_one(ChatHistory)
|
|
|
|
test_citations = [
|
|
Citation(
|
|
index=1,
|
|
document_id="doc1",
|
|
chunk_id="chunk1",
|
|
document_uri="file:///test/doc1.pdf",
|
|
document_title="Test Document 1",
|
|
page_numbers=[1, 2],
|
|
headings=["Section 1"],
|
|
content="This is some test content from doc 1",
|
|
),
|
|
Citation(
|
|
index=2,
|
|
document_id="doc2",
|
|
chunk_id="chunk2",
|
|
document_uri="file:///test/doc2.pdf",
|
|
document_title="Test Document 2",
|
|
page_numbers=[5],
|
|
headings=["Section 2", "Subsection"],
|
|
content="This is test content from doc 2",
|
|
),
|
|
]
|
|
|
|
await chat_history.add_citations(test_citations)
|
|
|
|
citation_widgets = chat_history.query(CitationWidget)
|
|
assert len(list(citation_widgets)) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_history_thinking_indicator(temp_db_path: Path):
|
|
"""Test that ChatHistory can show and hide thinking indicator."""
|
|
from haiku.rag.chat.widgets.chat_history import ChatHistory, ThinkingWidget
|
|
|
|
app, mock_client = _make_app(temp_db_path)
|
|
|
|
with (
|
|
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
|
|
_covering_returns(_stub_rag, mock_client),
|
|
):
|
|
async with app.run_test() as pilot:
|
|
chat_history = app.query_one(ChatHistory)
|
|
|
|
await chat_history.show_thinking()
|
|
thinking = chat_history.query(ThinkingWidget)
|
|
assert len(list(thinking)) == 1
|
|
|
|
chat_history.hide_thinking()
|
|
await pilot.pause()
|
|
thinking = chat_history.query(ThinkingWidget)
|
|
assert len(list(thinking)) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_clear_chat_resets_state(temp_db_path: Path):
|
|
"""Test that clearing chat resets state, messages, and conversation id."""
|
|
from haiku.rag.chat.widgets.chat_history import ChatHistory
|
|
|
|
app, mock_client = _make_app(temp_db_path)
|
|
|
|
with (
|
|
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
|
|
_covering_returns(_stub_rag, mock_client),
|
|
):
|
|
async with app.run_test() as pilot:
|
|
chat_history = app.query_one(ChatHistory)
|
|
|
|
await chat_history.add_message("user", "Hello")
|
|
await chat_history.add_message("assistant", "Hi there")
|
|
assert len(chat_history.messages) == 2
|
|
|
|
previous_conversation_id = app._conversation_id
|
|
await app.action_clear_chat()
|
|
await pilot.pause()
|
|
|
|
assert len(chat_history.messages) == 0
|
|
assert app._conversation_id != previous_conversation_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_citation_expand_collapse_with_enter(temp_db_path: Path):
|
|
"""Test that pressing Enter on a focused citation toggles expand/collapse."""
|
|
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
|
|
from haiku.rag.store.models.citation import Citation
|
|
|
|
app, mock_client = _make_app(temp_db_path)
|
|
|
|
with (
|
|
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
|
|
_covering_returns(_stub_rag, mock_client),
|
|
):
|
|
async with app.run_test() as pilot:
|
|
chat_history = app.query_one(ChatHistory)
|
|
|
|
test_citation = Citation(
|
|
index=1,
|
|
document_id="doc1",
|
|
chunk_id="chunk1",
|
|
document_uri="file:///test/doc1.pdf",
|
|
document_title="Test Document",
|
|
page_numbers=[1],
|
|
content="Test content",
|
|
)
|
|
await chat_history.add_citations([test_citation])
|
|
|
|
citation_widget = chat_history.query_one(CitationWidget)
|
|
assert citation_widget.collapsed is True
|
|
|
|
citation_widget.focus()
|
|
await pilot.pause()
|
|
|
|
await pilot.press("enter")
|
|
await pilot.pause()
|
|
assert citation_widget.collapsed is False
|
|
|
|
await pilot.press("enter")
|
|
await pilot.pause()
|
|
assert citation_widget.collapsed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_show_citations_renders_from_flat_state(temp_db_path: Path):
|
|
"""Citations in state (flat list[str]) render into the chat history."""
|
|
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
|
|
|
|
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() as pilot:
|
|
rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE])
|
|
|
|
citation = Citation(
|
|
index=1,
|
|
document_id="doc1",
|
|
chunk_id="chunk1",
|
|
document_uri="file:///test/doc1.pdf",
|
|
document_title="Test Document",
|
|
page_numbers=[1],
|
|
content="Cited content",
|
|
)
|
|
rag_state.citation_index["chunk1"] = citation
|
|
rag_state.citations.append("chunk1")
|
|
app._state[RAG_STATE_NAMESPACE] = rag_state.model_dump(mode="json")
|
|
|
|
chat_history = app.query_one(ChatHistory)
|
|
await app._show_citations_and_programs(chat_history)
|
|
await pilot.pause()
|
|
|
|
widgets = list(chat_history.query(CitationWidget))
|
|
assert len(widgets) == 1
|
|
assert widgets[0].citation.chunk_id == "chunk1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_document_filter_updates_rag_state(temp_db_path: Path):
|
|
"""Test that selecting document filters updates RAGState.document_filter."""
|
|
from haiku.rag.chat.app import RAG_STATE_NAMESPACE
|
|
from haiku.rag.chat.widgets.document_filter_modal import DocumentFilterModal
|
|
from haiku.rag.tools.filters import build_document_id_filter
|
|
|
|
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():
|
|
# The selection is document ids, so a repeated title cannot widen it.
|
|
selected = [
|
|
"6f1c2d4e-0000-4000-8000-000000000001",
|
|
"6f1c2d4e-0000-4000-8000-000000000002",
|
|
]
|
|
app.on_document_filter_modal_filter_changed(
|
|
DocumentFilterModal.FilterChanged(selected)
|
|
)
|
|
|
|
# RAGState.document_filter should be set
|
|
rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE])
|
|
expected_filter = build_document_id_filter(selected)
|
|
assert rag_state.document_filter == expected_filter
|
|
assert rag_state.document_filter is not None
|
|
assert "LIKE" not in rag_state.document_filter
|
|
|
|
# The state snapshot should also reflect the change
|
|
assert app._state["rag"]["document_filter"] == expected_filter
|
|
|
|
|
|
@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."""
|
|
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():
|
|
# First set a filter
|
|
app.on_document_filter_modal_filter_changed(
|
|
DocumentFilterModal.FilterChanged(["AI Overview"])
|
|
)
|
|
rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE])
|
|
assert rag_state.document_filter is not None
|
|
|
|
# Then clear it
|
|
app.on_document_filter_modal_filter_changed(
|
|
DocumentFilterModal.FilterChanged([])
|
|
)
|
|
rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE])
|
|
assert rag_state.document_filter is None
|
|
assert app._state["rag"]["document_filter"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_app_open_failure_surfaces_real_error(tmp_path: Path):
|
|
"""A failed database open must surface its own error, not an
|
|
AttributeError from tearing down a client that never opened."""
|
|
from haiku.rag.chat.app import ChatApp
|
|
|
|
app = ChatApp(scope=for_path(tmp_path / "missing.lancedb"), capabilities=[])
|
|
with pytest.raises(FileNotFoundError):
|
|
async with app.run_test():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_cancelled_run_does_not_advance_persisted_state(temp_db_path: Path):
|
|
"""State and message history have to move together, or the thread bricks.
|
|
|
|
A cancelled run keeps whatever the tools wrote but discards the run's messages.
|
|
If the state advanced, the next question derives its identity from the shorter
|
|
history, lands behind the recorded evidence epoch, and is refused as
|
|
non-append-only — leaving the conversation unusable until cleared.
|
|
"""
|
|
import asyncio
|
|
|
|
app, mock_client = _make_app(temp_db_path)
|
|
|
|
class CancellingRun:
|
|
"""A run that writes evidence through the tools, then is cancelled."""
|
|
|
|
def __init__(self, deps):
|
|
self._deps = deps
|
|
|
|
async def __aenter__(self):
|
|
self._deps.state["rag"] = {
|
|
"evidence": {"question": 0, "latest_evidence_epoch": 7}
|
|
}
|
|
return self
|
|
|
|
async def __aexit__(self, *_):
|
|
return False
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
raise asyncio.CancelledError
|
|
|
|
with (
|
|
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
|
|
_covering_returns(_stub_rag, mock_client),
|
|
):
|
|
async with app.run_test():
|
|
app._state = {
|
|
"rag": {"evidence": {"question": 0, "latest_evidence_epoch": 0}}
|
|
}
|
|
before = deepcopy(app._state)
|
|
|
|
class Agent:
|
|
def run_stream_events(self, *_, deps, **__):
|
|
return CancellingRun(deps)
|
|
|
|
app._agent = Agent() # type: ignore[assignment]
|
|
await app._run_agent("a question")
|
|
|
|
assert app._state == before
|
|
assert app._messages == []
|
|
|
|
|
|
async def test_visual_grounding_uses_the_database_holding_the_citation(tmp_path):
|
|
"""Chunks, pages and boxes come from one database. Covering a set, the
|
|
citation's source says which, and a covering client has no repositories."""
|
|
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
|
|
from haiku.rag.store.models import Chunk
|
|
from haiku.rag.store.models.citation import Citation
|
|
|
|
owner = _make_mock_client()
|
|
owner.source = "beta"
|
|
owner.get_chunk_by_id.return_value = Chunk(
|
|
id="c1", document_id="d1", content="cited body"
|
|
)
|
|
|
|
covering = _make_mock_client()
|
|
covering.covers_multiple = True
|
|
covering.source_names = ("alpha", "beta")
|
|
covering.reader_for = AsyncMock(return_value=owner)
|
|
|
|
app, _ = _make_app(tmp_path / "unused.lancedb", covering)
|
|
with (
|
|
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
|
|
_covering_returns(_stub_rag, covering),
|
|
):
|
|
async with app.run_test():
|
|
history = app.query_one(ChatHistory)
|
|
await history.add_citations(
|
|
[
|
|
Citation(
|
|
chunk_id="c1",
|
|
document_id="d1",
|
|
document_uri="test://beta/one",
|
|
content="cited body",
|
|
source="beta",
|
|
)
|
|
]
|
|
)
|
|
widget = next(iter(app.query(CitationWidget)))
|
|
widget.add_class("selected")
|
|
|
|
with patch.object(app, "push_screen", new=AsyncMock()) as push:
|
|
await app.action_show_visual()
|
|
|
|
covering.reader_for.assert_awaited_once_with("beta")
|
|
owner.get_chunk_by_id.assert_awaited_once_with("c1")
|
|
assert push.await_args is not None
|
|
assert push.await_args.args[0].client is owner
|
|
|
|
|
|
class TestLendingTheClient:
|
|
@pytest.mark.asyncio
|
|
async def test_mounting_lends_its_client_to_every_capability(self, temp_db_path):
|
|
"""Capabilities are built before the client exists, so each reads through
|
|
the one the app opened rather than opening its own."""
|
|
client = _make_mock_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():
|
|
borrowed = [c.borrowed_rag for c in app._capabilities]
|
|
|
|
assert borrowed == [client] * len(app._capabilities)
|
|
assert borrowed
|
|
|
|
|
|
class TestDocumentSelectionIdentity:
|
|
"""Two documents can share a title, within a corpus and across databases, so
|
|
the selection is by id and the label says which database."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_repeated_title_selects_one_document(self, temp_db_path: Path):
|
|
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 = ("arxiv", "wiki")
|
|
client.list_documents.return_value = [
|
|
Document(id="id-one", content="", title="Capital region", source="arxiv"),
|
|
Document(id="id-two", content="", title="Capital region", source="wiki"),
|
|
]
|
|
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))
|
|
selected_ids = [box.doc_id for box in boxes]
|
|
assert selected_ids == ["id-one", "id-two"]
|
|
labels = [str(b.label) for b in boxes]
|
|
assert "Capital region (arxiv)" in labels
|
|
assert "Capital region (wiki)" in labels
|
|
|
|
boxes[0].value = True
|
|
await pilot.pause()
|
|
assert modal._selected == {"id-one"}
|
|
|
|
|
|
class TestDocumentSearchFilter:
|
|
"""The filter modal shows one page and asks the database for the rest, so the
|
|
typed term reaches SQL."""
|
|
|
|
def test_no_term_means_no_filter(self):
|
|
from haiku.rag.chat.widgets.document_filter_modal import search_filter
|
|
|
|
assert search_filter(" ") is None
|
|
|
|
def test_a_term_matches_titles_and_uris(self):
|
|
from haiku.rag.chat.widgets.document_filter_modal import search_filter
|
|
|
|
built = search_filter("Nobel")
|
|
|
|
assert built == (
|
|
"LOWER(title) LIKE '%nobel%' ESCAPE '\\' "
|
|
"OR LOWER(uri) LIKE '%nobel%' ESCAPE '\\'"
|
|
)
|
|
|
|
def test_the_search_is_case_insensitive(self):
|
|
"""LIKE is case-sensitive, so both sides are lowered."""
|
|
from haiku.rag.chat.widgets.document_filter_modal import search_filter
|
|
|
|
built = search_filter("NoBeL")
|
|
|
|
assert built is not None
|
|
assert "LOWER(title) LIKE '%nobel%'" in built
|
|
assert "LOWER(uri) LIKE '%nobel%'" in built
|
|
|
|
def test_like_wildcards_in_the_term_are_literal(self):
|
|
"""A term is text to find, not a pattern: `_` matches any character."""
|
|
from haiku.rag.chat.widgets.document_filter_modal import search_filter
|
|
|
|
built = search_filter("100%_raw")
|
|
|
|
assert built is not None
|
|
assert "100\\%\\_raw" in built
|
|
assert "ESCAPE '\\'" in built
|
|
|
|
def test_a_quote_in_the_term_is_escaped(self):
|
|
"""Whatever was typed is data, not syntax."""
|
|
from haiku.rag.chat.widgets.document_filter_modal import search_filter
|
|
|
|
built = search_filter("O'Brien")
|
|
|
|
assert built is not None
|
|
assert "o''brien" in built
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submitting_a_search_reloads_the_page(self, temp_db_path: Path):
|
|
"""The typed term reaches the database and replaces what is shown."""
|
|
from textual.widgets import Input
|
|
|
|
from haiku.rag.chat.widgets.document_filter_modal import (
|
|
DocumentCheckbox,
|
|
DocumentFilterModal,
|
|
)
|
|
from haiku.rag.store.models.document import Document
|
|
|
|
client = AsyncMock()
|
|
client.covers_multiple = False
|
|
client.source_names = ()
|
|
client.list_documents.return_value = [
|
|
Document(id="id-one", content="", title="Capital region"),
|
|
Document(id="id-two", content="", title="Nobel laureates"),
|
|
]
|
|
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()
|
|
assert len(list(modal.query(DocumentCheckbox))) == 2
|
|
|
|
client.list_documents.return_value = [
|
|
Document(id="id-two", content="", title="Nobel laureates"),
|
|
]
|
|
client.count_documents.return_value = 1
|
|
await modal.on_input_submitted(Input.Submitted(Input(), "Nobel"))
|
|
await pilot.pause()
|
|
|
|
assert client.list_documents.await_args is not None
|
|
assert "nobel" in client.list_documents.await_args.kwargs["filter"]
|
|
labels = [str(b.label) for b in modal.query(DocumentCheckbox)]
|
|
assert labels == ["Nobel laureates"]
|