CLI app
This commit is contained in:
parent
100933b36e
commit
6b3ff5ef1e
8 changed files with 1059 additions and 0 deletions
|
|
@ -16,6 +16,11 @@
|
||||||
- `CitationInfo` and `QAResponse` models for structured responses
|
- `CitationInfo` and `QAResponse` models for structured responses
|
||||||
- Natural language document filtering via `build_document_filter()`
|
- Natural language document filtering via `build_document_filter()`
|
||||||
- Configurable search limit per agent
|
- 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
|
- **Q/A History Management**: Intelligent conversation history with semantic ranking
|
||||||
- FIFO queue with 50 max entries
|
- FIFO queue with 50 max entries
|
||||||
- Embedding cache to avoid re-embedding Q/A pairs
|
- Embedding cache to avoid re-embedding Q/A pairs
|
||||||
|
|
|
||||||
31
haiku_rag_slim/haiku/rag/chat/__init__.py
Normal file
31
haiku_rag_slim/haiku/rag/chat/__init__.py
Normal file
|
|
@ -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()
|
||||||
324
haiku_rag_slim/haiku/rag/chat/app.py
Normal file
324
haiku_rag_slim/haiku/rag/chat/app.py
Normal file
|
|
@ -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
|
||||||
0
haiku_rag_slim/haiku/rag/chat/widgets/__init__.py
Normal file
0
haiku_rag_slim/haiku/rag/chat/widgets/__init__.py
Normal file
354
haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py
Normal file
354
haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py
Normal file
|
|
@ -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()
|
||||||
|
|
@ -540,6 +540,21 @@ def inspect(
|
||||||
run_inspector(db_path, read_only=_read_only, before=_before)
|
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(
|
@cli.command(
|
||||||
"serve",
|
"serve",
|
||||||
help="Start haiku.rag server. Use --monitor and/or --mcp to enable services.",
|
help="Start haiku.rag server. Use --monitor and/or --mcp to enable services.",
|
||||||
|
|
|
||||||
0
tests/chat/__init__.py
Normal file
0
tests/chat/__init__.py
Normal file
330
tests/chat/test_chat_app.py
Normal file
330
tests/chat/test_chat_app.py
Normal file
|
|
@ -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
|
||||||
Loading…
Reference in a new issue