diff --git a/CHANGELOG.md b/CHANGELOG.md index 2496c1d5..54803dfd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,11 @@ - `CitationInfo` and `QAResponse` models for structured responses - Natural language document filtering via `build_document_filter()` - Configurable search limit per agent +- **Chat TUI** (`haiku-rag chat`): Terminal-based chat interface using Textual + - Single chat window with inline tool calls and expandable citations + - Visual grounding (`v` key) reuses inspector's `VisualGroundingModal` + - Database info (`i` key) shows document/chunk counts and storage info + - Keybindings: `q` quit, `Ctrl+L` clear chat, `Escape` focus input - **Q/A History Management**: Intelligent conversation history with semantic ranking - FIFO queue with 50 max entries - Embedding cache to avoid re-embedding Q/A pairs diff --git a/haiku_rag_slim/haiku/rag/chat/__init__.py b/haiku_rag_slim/haiku/rag/chat/__init__.py new file mode 100644 index 00000000..73fc5547 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/chat/__init__.py @@ -0,0 +1,31 @@ +from datetime import datetime +from pathlib import Path + + +def run_chat( + db_path: Path | None = None, + read_only: bool = False, + before: datetime | None = None, +) -> None: + """Run the chat TUI. + + Args: + db_path: Path to the LanceDB database. If None, uses default from config. + read_only: Whether to open the database in read-only mode. + before: Query database as it existed before this datetime. + """ + try: + from haiku.rag.chat.app import ChatApp + except ImportError as e: + raise ImportError( + "Textual is required for the chat TUI. Install with: pip install textual" + ) from e + + from haiku.rag.config import get_config + + config = get_config() + if db_path is None: + db_path = config.storage.data_dir / "haiku.rag.lancedb" + + app = ChatApp(db_path, read_only=read_only, before=before) + app.run() diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py new file mode 100644 index 00000000..7536ac96 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -0,0 +1,324 @@ +# pyright: reportPossiblyUnboundVariable=false +import asyncio +import uuid +from collections.abc import AsyncIterable +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from ag_ui.core import EventType +from pydantic_ai import ( + Agent, + AgentStreamEvent, + FunctionToolCallEvent, + FunctionToolResultEvent, + RunContext, +) +from pydantic_ai.messages import ModelMessage + +from haiku.rag.agents.chat.agent import create_chat_agent +from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState, CitationInfo +from haiku.rag.client import HaikuRAG +from haiku.rag.config import get_config + +if TYPE_CHECKING: + from textual.app import ComposeResult + +try: + import logfire + + logfire.configure(console=False) +except ImportError: + pass + +try: + import textual_image.widget # noqa: F401 - import early for renderer detection + from textual.app import App + from textual.binding import Binding + from textual.widgets import Footer, Header, Input + from textual.worker import Worker + + from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget + + TEXTUAL_AVAILABLE = True +except ImportError: + TEXTUAL_AVAILABLE = False + App = object # type: ignore + + +class ChatApp(App): # type: ignore[misc] + """Textual TUI for conversational RAG.""" + + TITLE = "haiku.rag Chat" + + CSS = """ + Screen { + layout: grid; + grid-size: 1 2; + grid-rows: 1fr auto; + background: $surface; + } + + #chat-history { + height: 100%; + } + + Header { + background: $primary; + } + + Footer { + background: $surface-darken-1; + } + """ + + BINDINGS = [ + Binding("ctrl+l", "clear_chat", "Clear", show=True), + Binding("ctrl+g", "show_visual", "Visual", show=True), + Binding("ctrl+i", "show_info", "Info", show=True), + Binding("escape", "focus_input", "Focus Input", show=False), + ] + + def __init__( + self, db_path: Path, read_only: bool = False, before: datetime | None = None + ) -> None: + super().__init__() + self.db_path = db_path + self.read_only = read_only + self.before = before + self.client: HaikuRAG | None = None + self.config = get_config() + self.agent: Agent[ChatDeps, str] | None = None + self.session_state: ChatSessionState | None = None + self._is_processing = False + self._tool_call_widgets: dict[str, Any] = {} + self._last_citations: list[CitationInfo] = [] + self._selected_citation_idx: int | None = None + self._current_worker: Worker[None] | None = None + self._message_history: list[ModelMessage] = [] + + def compose(self) -> "ComposeResult": + """Compose the UI layout.""" + yield Header() + yield ChatHistory(id="chat-history") + yield Input(placeholder="Ask a question...", id="chat-input") + yield Footer() + + async def on_mount(self) -> None: + """Initialize the app when mounted.""" + self.client = HaikuRAG( + db_path=self.db_path, + config=self.config, + read_only=self.read_only, + before=self.before, + ) + await self.client.__aenter__() + + # Create agent and session state + self.agent = create_chat_agent(self.config) + self.session_state = ChatSessionState(session_id=str(uuid.uuid4())) + + # Focus the input field + self.query_one(Input).focus() + + async def on_unmount(self) -> None: + """Clean up when unmounting.""" + if self.client: + await self.client.__aexit__(None, None, None) + + async def _handle_stream_event(self, event: AgentStreamEvent) -> None: + """Handle streaming events from the agent.""" + chat_history = self.query_one(ChatHistory) + + if isinstance(event, FunctionToolCallEvent): + tool_name = event.part.tool_name + tool_call_id = event.part.tool_call_id or str(uuid.uuid4()) + args = event.part.args_as_dict() + widget = await chat_history.add_tool_call(tool_name, args) + self._tool_call_widgets[tool_call_id] = widget + + elif isinstance(event, FunctionToolResultEvent): + tool_call_id = event.tool_call_id + if tool_call_id and tool_call_id in self._tool_call_widgets: + widget = self._tool_call_widgets[tool_call_id] + chat_history.mark_tool_complete(widget) + + # Extract citations from StateSnapshotEvent in tool metadata + result = getattr(event, "result", None) + metadata = getattr(result, "metadata", None) if result else None + if metadata: + for meta_event in metadata: + if ( + hasattr(meta_event, "type") + and meta_event.type == EventType.STATE_SNAPSHOT + ): + snapshot = getattr(meta_event, "snapshot", {}) + if "citations" in snapshot: + self._last_citations = [ + CitationInfo(**c) for c in snapshot["citations"] + ] + + async def _event_stream_handler( + self, + _ctx: RunContext[ChatDeps], + event_stream: AsyncIterable[AgentStreamEvent], + ) -> None: + """Handle streaming events from the agent.""" + async for event in event_stream: + await self._handle_stream_event(event) + # Yield to event loop to keep UI responsive + await asyncio.sleep(0) + + async def on_input_submitted(self, event: Input.Submitted) -> None: + """Handle user input submission.""" + user_message = event.value.strip() + if not user_message or self._is_processing: + return + + if not self.client or not self.agent: + return + + # Clear the input + event.input.clear() + + # Add user message to history + chat_history = self.query_one(ChatHistory) + await chat_history.add_message("user", user_message) + + # Clear for new query + self._tool_call_widgets.clear() + self._last_citations.clear() + self._selected_citation_idx = None + + # Run agent in a worker to keep UI responsive + self._is_processing = True + self.query_one(Input).disabled = True + self._current_worker = self.run_worker( + self._run_agent(user_message), exclusive=True + ) + + async def _run_agent(self, user_message: str) -> None: + """Run the agent in a background worker.""" + if not self.client or not self.agent: + return + + chat_history = self.query_one(ChatHistory) + + # Show thinking indicator + await chat_history.show_thinking() + + try: + deps = ChatDeps( + client=self.client, + config=self.config, + session_state=self.session_state, + ) + + async with self.agent.run_stream( + user_message, + deps=deps, + message_history=self._message_history, + event_stream_handler=self._event_stream_handler, + ) as stream: + # Hide thinking when we start getting content + chat_history.hide_thinking() + + # Create assistant message for streaming + assistant_msg = await chat_history.add_message("assistant", "") + + # Stream text updates + async for text in stream.stream_text(): + assistant_msg.update_content(text) + chat_history.scroll_end(animate=False) + # Yield to event loop to keep UI responsive + await asyncio.sleep(0) + + # Update message history with this conversation + self._message_history = stream.all_messages() + + # Add citations captured from tool metadata + if self._last_citations: + await chat_history.add_citations(self._last_citations) + + except asyncio.CancelledError: + chat_history.hide_thinking() + await chat_history.add_message("assistant", "*Cancelled*") + except Exception as e: + chat_history.hide_thinking() + await chat_history.add_message("assistant", f"Error: {e}") + finally: + self._is_processing = False + self._current_worker = None + chat_input = self.query_one(Input) + chat_input.disabled = False + chat_input.focus() + + async def action_clear_chat(self) -> None: + """Clear the chat history and reset session.""" + chat_history = self.query_one(ChatHistory) + await chat_history.clear_messages() + self._last_citations.clear() + self._selected_citation_idx = None + self._message_history.clear() + # Reset session state for fresh conversation + self.session_state = ChatSessionState(session_id=str(uuid.uuid4())) + + def action_focus_input(self) -> None: + """Focus the input field, or cancel if processing.""" + if self._is_processing and self._current_worker: + self._current_worker.cancel() + self.query_one(Input).focus() + + def _clear_citation_selection(self) -> None: + """Clear citation selection.""" + chat_history = self.query_one(ChatHistory) + for widget in chat_history.query(CitationWidget): + widget.remove_class("selected") + self._selected_citation_idx = None + + def on_descendant_focus(self, _event: object) -> None: + """Clear citation selection when input is focused.""" + if isinstance(self.focused, Input): + self._clear_citation_selection() + + async def action_show_visual(self) -> None: + """Show visual grounding for the selected citation.""" + if not self.client or not self._last_citations: + return + + idx = ( + self._selected_citation_idx + if self._selected_citation_idx is not None + else 0 + ) + citation = self._last_citations[idx] + chunk = await self.client.chunk_repository.get_by_id(citation.chunk_id) + if not chunk: + return + + from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal + + await self.push_screen(VisualGroundingModal(chunk=chunk, client=self.client)) + + async def action_show_info(self) -> None: + """Show database info modal.""" + if not self.client: + return + + from haiku.rag.inspector.widgets.info_modal import InfoModal + + await self.push_screen(InfoModal(self.client, self.db_path)) + + def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None: + """Handle citation selection.""" + chat_history = self.query_one(ChatHistory) + + # Remove selected class from all citations + 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") + self._selected_citation_idx = event.citation_index diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/__init__.py b/haiku_rag_slim/haiku/rag/chat/widgets/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py b/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py new file mode 100644 index 00000000..d6b4b0c9 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py @@ -0,0 +1,354 @@ +from typing import TYPE_CHECKING + +from textual.containers import Horizontal, VerticalScroll +from textual.message import Message +from textual.widgets import Collapsible, LoadingIndicator, Markdown, Static + +from haiku.rag.agents.chat.state import CitationInfo + +if TYPE_CHECKING: + from textual.app import ComposeResult + from textual.events import Key + + +class ChatMessage(Static): + """A single chat message with role styling.""" + + def __init__(self, role: str, content: str = "", **kwargs) -> None: + super().__init__(**kwargs) + self.role = role + self.content = content + + def compose(self) -> "ComposeResult": + prefix = "**You:**" if self.role == "user" else "**Assistant:**" + yield Markdown(f"{prefix}\n\n{self.content}", id="message-content") + + def update_content(self, content: str) -> None: + """Update the message content (for streaming).""" + self.content = content + prefix = "**You:**" if self.role == "user" else "**Assistant:**" + markdown = self.query_one("#message-content", Markdown) + markdown.update(f"{prefix}\n\n{content}") + + +class ToolCallWidget(Static): + """Styled inline display of a tool call.""" + + TOOL_LABELS = { + "search": "Searching", + "ask": "Asking", + "get_document": "Fetching", + } + + def __init__(self, tool_name: str, args: dict | None = None, **kwargs) -> None: + super().__init__(**kwargs) + self.tool_name = tool_name + self.args = args or {} + self._complete = False + + def compose(self) -> "ComposeResult": + label = self.TOOL_LABELS.get(self.tool_name, self.tool_name) + + # Build description based on tool type + if self.tool_name == "search": + query = self.args.get("query", "...") + doc = self.args.get("document_name") + desc = f'"{query}"' + if doc: + desc += f" in {doc}" + elif self.tool_name == "ask": + question = self.args.get("question", "...") + doc = self.args.get("document_name") + desc = f'"{question}"' + if doc: + desc += f" from {doc}" + elif self.tool_name == "get_document": + query = self.args.get("query", "...") + desc = f'"{query}"' + else: + desc = str(self.args) if self.args else "" + + with Horizontal(classes="tool-row"): + if self._complete: + yield Static("✓", classes="tool-status") + else: + yield LoadingIndicator(classes="tool-spinner") + yield Static(label, classes="tool-badge") + yield Static(desc, classes="tool-desc") + + def mark_complete(self) -> None: + self._complete = True + self.refresh(recompose=True) + + +class CitationWidget(Collapsible): + """Inline expandable citation.""" + + can_focus = True + can_focus_children = False + + class Selected(Message): + """Message sent when a citation is selected.""" + + def __init__(self, citation_index: int) -> None: + super().__init__() + self.citation_index = citation_index + + def __init__(self, citation: CitationInfo, **kwargs) -> None: + title = f"[{citation.index}] {citation.document_title or citation.document_uri}" + if citation.page_numbers: + pages = ", ".join(map(str, citation.page_numbers[:3])) + if len(citation.page_numbers) > 3: + pages += "..." + title += f" (p.{pages})" + + # Build content widgets + content = citation.content + if len(content) > 500: + content = content[:500] + "..." + + children: list[Markdown | Static] = [Markdown(content)] + if citation.headings: + headings = " > ".join(citation.headings[:3]) + children.append(Static(f"Section: {headings}", classes="citation-metadata")) + children.append( + Static(f"Source: {citation.document_uri}", classes="citation-metadata") + ) + + super().__init__(*children, title=title, collapsed=True, **kwargs) + self.citation = citation + + def on_focus(self) -> None: + """When focused, mark as selected.""" + self.post_message(self.Selected(self.citation.index - 1)) + + def on_key(self, event: "Key") -> None: + """Handle Enter to toggle expand/collapse.""" + if event.key == "enter": + self.collapsed = not self.collapsed + event.stop() + + +class ThinkingWidget(Static): + """Thinking indicator shown while agent is processing.""" + + def compose(self) -> "ComposeResult": + with Horizontal(classes="thinking-row"): + yield LoadingIndicator(classes="thinking-spinner") + yield Static("Thinking...", classes="thinking-text") + + +class SourcesHeader(Static): + """Header for the citations section.""" + + def __init__(self, count: int, **kwargs) -> None: + super().__init__(f"Sources ({count})", **kwargs) + + +class ChatHistory(VerticalScroll): + """Scrollable container for chat messages, tool calls, and citations.""" + + can_focus = True + + DEFAULT_CSS = """ + ChatHistory { + height: 100%; + background: $surface; + padding: 1 2; + } + + /* Messages */ + ChatMessage { + margin: 1 0; + padding: 1 2; + background: $panel; + } + + ChatMessage.user { + background: $primary 15%; + border-left: thick $primary; + margin-right: 4; + } + + ChatMessage.assistant { + background: $success 15%; + border-left: thick $success; + margin-left: 4; + } + + ChatMessage Markdown { + margin: 0; + padding: 0; + } + + /* Tool calls */ + ToolCallWidget { + margin: 0 0 0 4; + padding: 0 1; + height: auto; + background: $surface; + border-left: thick $warning; + } + + ToolCallWidget.complete { + border-left: thick $success; + } + + .tool-row { + height: auto; + width: 100%; + } + + .tool-spinner { + width: 2; + height: 1; + color: $warning; + } + + .tool-status { + width: 2; + color: $success; + } + + .tool-badge { + width: auto; + color: $text; + text-style: bold; + padding-right: 1; + } + + .tool-desc { + width: 1fr; + color: $text-muted; + } + + /* Sources section */ + SourcesHeader { + margin: 2 0 1 0; + padding: 0 1; + text-style: bold; + color: $text; + background: $primary 15%; + border-left: thick $primary; + } + + /* Citations */ + CitationWidget { + margin: 0 0 0 2; + background: $surface; + } + + CitationWidget > CollapsibleTitle { + padding: 0 1; + color: $text-muted; + } + + CitationWidget:focus { + background: $accent 15%; + border-left: thick $accent; + } + + CitationWidget:focus > CollapsibleTitle { + color: $text; + text-style: bold; + } + + CitationWidget.selected { + background: $accent 15%; + border-left: thick $accent; + } + + CitationWidget.selected > CollapsibleTitle { + color: $text; + text-style: bold; + } + + CitationWidget Contents { + padding: 1 2; + background: $panel; + } + + CitationWidget .citation-metadata { + margin-top: 1; + color: $text-muted; + text-style: italic; + } + + /* Thinking indicator */ + ThinkingWidget { + margin: 1 0 0 4; + padding: 0 1; + height: auto; + background: $surface; + border-left: thick $primary; + } + + .thinking-row { + height: auto; + width: 100%; + } + + .thinking-spinner { + width: 2; + height: 1; + color: $primary; + } + + .thinking-text { + color: $text-muted; + text-style: italic; + } + """ + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.messages: list[tuple[str, str]] = [] + + async def add_message(self, role: str, content: str = "") -> ChatMessage: + """Add a message to the chat history.""" + self.messages.append((role, content)) + message_widget = ChatMessage(role, content, classes=role) + await self.mount(message_widget) + self.scroll_end(animate=False) + return message_widget + + async def add_tool_call( + self, tool_name: str, args: dict | None = None + ) -> ToolCallWidget: + """Add an inline tool call indicator.""" + widget = ToolCallWidget(tool_name, args) + await self.mount(widget) + self.scroll_end(animate=False) + return widget + + def mark_tool_complete(self, widget: ToolCallWidget) -> None: + """Mark a tool call as complete.""" + widget.mark_complete() + widget.add_class("complete") + + async def add_citations(self, citations: list[CitationInfo]) -> None: + """Add citations inline after a response.""" + if not citations: + return + await self.mount(SourcesHeader(len(citations))) + for citation in citations: + widget = CitationWidget(citation) + await self.mount(widget) + self.scroll_end(animate=False) + + async def show_thinking(self) -> None: + """Show the thinking indicator.""" + await self.mount(ThinkingWidget(id="thinking")) + self.scroll_end(animate=False) + + def hide_thinking(self) -> None: + """Hide the thinking indicator.""" + try: + self.query_one("#thinking", ThinkingWidget).remove() + except Exception: + pass + + async def clear_messages(self) -> None: + """Clear all messages from the chat history.""" + self.messages.clear() + await self.remove_children() diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 9958cc07..11649c17 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -540,6 +540,21 @@ def inspect( run_inspector(db_path, read_only=_read_only, before=_before) +@cli.command("chat", help="Launch interactive chat TUI for conversational RAG") +def chat( + db: Path | None = typer.Option( + None, + "--db", + help="Path to the LanceDB database file", + ), +): + """Launch the chat TUI for conversational RAG.""" + from haiku.rag.chat import run_chat + + db_path = db if db else get_config().storage.data_dir / "haiku.rag.lancedb" + run_chat(db_path, read_only=_read_only, before=_before) + + @cli.command( "serve", help="Start haiku.rag server. Use --monitor and/or --mcp to enable services.", diff --git a/tests/chat/__init__.py b/tests/chat/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py new file mode 100644 index 00000000..66e7f7e4 --- /dev/null +++ b/tests/chat/test_chat_app.py @@ -0,0 +1,330 @@ +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from typer.testing import CliRunner + +from haiku.rag.cli import cli + +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() + + +@pytest.mark.asyncio +async def test_chat_app_has_required_widgets(temp_db_path: Path): + """Test that ChatApp has the required widgets: ChatHistory, Input.""" + from haiku.rag.chat.app import ChatApp + from haiku.rag.chat.widgets.chat_history import ChatHistory + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test(): + chat_history = app.query_one(ChatHistory) + assert chat_history is not None + + from textual.widgets import Input + + chat_input = app.query_one(Input) + 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.""" + from haiku.rag.chat.app import ChatApp + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test() as pilot: + # App should be running + assert app.is_running + + # Press ctrl+q to quit + await pilot.press("ctrl+q") + + # App should have exited + 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.app import ChatApp + from haiku.rag.chat.widgets.chat_history import ChatHistory + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test(): + chat_history = app.query_one(ChatHistory) + + # Add a user message + 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?") + + # Add an assistant message + 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_app_calls_agent_on_submit(temp_db_path: Path): + """Test that submitting a message triggers agent invocation.""" + from contextlib import asynccontextmanager + + from haiku.rag.chat.app import ChatApp + from haiku.rag.chat.widgets.chat_history import ChatHistory + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + # Create a mock stream result + mock_result = MagicMock() + mock_result.output = "This is the agent's response." + + # Mock stream that returns the result + mock_stream = MagicMock() + mock_stream.get_result = AsyncMock(return_value=mock_result) + + async def mock_stream_text(): + yield "This is the agent's response." + + mock_stream.stream_text = mock_stream_text + + @asynccontextmanager + async def mock_run_stream(*args, **kwargs): + yield mock_stream + + # Create a mock agent that returns a fixed response + mock_agent = MagicMock() + mock_agent.run_stream = mock_run_stream + + with ( + patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client), + patch("haiku.rag.chat.app.create_chat_agent", return_value=mock_agent), + ): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test() as pilot: + # Type a message in the input + await pilot.press("H", "e", "l", "l", "o") + await pilot.press("enter") + + # Give the app time to process + await pilot.pause() + + # Verify the response was added to chat history + chat_history = app.query_one(ChatHistory) + assert len(chat_history.messages) >= 2 # User message + agent response + + +@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.app import ChatApp + from haiku.rag.chat.widgets.chat_history import ChatHistory, ToolCallWidget + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test(): + chat_history = app.query_one(ChatHistory) + + # Add a tool call + tool_widget = await chat_history.add_tool_call("search", {"query": "test"}) + assert isinstance(tool_widget, ToolCallWidget) + assert tool_widget._complete is False + + # Mark it complete + chat_history.mark_tool_complete(tool_widget) + assert tool_widget._complete 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.agents.chat.state import CitationInfo + from haiku.rag.chat.app import ChatApp + from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test(): + chat_history = app.query_one(ChatHistory) + + test_citations = [ + CitationInfo( + 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", + ), + CitationInfo( + 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) + + # Verify citation widgets were added + 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.app import ChatApp + from haiku.rag.chat.widgets.chat_history import ChatHistory, ThinkingWidget + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test() as pilot: + chat_history = app.query_one(ChatHistory) + + # Show thinking indicator + await chat_history.show_thinking() + thinking = chat_history.query(ThinkingWidget) + assert len(list(thinking)) == 1 + + # Hide thinking indicator (remove is deferred, so pause) + 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_session(temp_db_path: Path): + """Test that clearing chat resets the session state.""" + from haiku.rag.chat.app import ChatApp + from haiku.rag.chat.widgets.chat_history import ChatHistory + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test() as pilot: + chat_history = app.query_one(ChatHistory) + + # Add some messages + await chat_history.add_message("user", "Hello") + await chat_history.add_message("assistant", "Hi there") + assert len(chat_history.messages) == 2 + + # Store original session ID + assert app.session_state is not None + original_session_id = app.session_state.session_id + + # Clear chat + await pilot.press("ctrl+l") + await pilot.pause() + + # Verify messages cleared + assert len(chat_history.messages) == 0 + + # Verify session state reset (new session ID) + assert app.session_state is not None + assert app.session_state.session_id != original_session_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.agents.chat.state import CitationInfo + from haiku.rag.chat.app import ChatApp + from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + app = ChatApp(temp_db_path, read_only=True) + + async with app.run_test() as pilot: + chat_history = app.query_one(ChatHistory) + + # Add a citation + test_citation = CitationInfo( + 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]) + + # Get the citation widget + citation_widget = chat_history.query_one(CitationWidget) + assert citation_widget.collapsed is True + + # Focus the citation + citation_widget.focus() + await pilot.pause() + + # Press Enter to expand + await pilot.press("enter") + await pilot.pause() + assert citation_widget.collapsed is False + + # Press Enter again to collapse + await pilot.press("enter") + await pilot.pause() + assert citation_widget.collapsed is True