diff --git a/app/backend/main.py b/app/backend/main.py index 7bbcf4cf..5d9f0067 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -3,8 +3,9 @@ import os from pathlib import Path from dotenv import find_dotenv, load_dotenv +from pydantic_ai import Agent +from pydantic_ai.ag_ui import AGUIAdapter from pydantic_ai.ui import SSE_CONTENT_TYPE -from pydantic_ai.ui.ag_ui import AGUIAdapter from starlette.applications import Starlette from starlette.middleware import Middleware from starlette.middleware.cors import CORSMiddleware @@ -12,22 +13,14 @@ from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route -from haiku.rag.agents.chat import ( - AGUI_STATE_KEY, - ChatDeps, - build_chat_toolkit, - create_chat_agent, -) from haiku.rag.client import HaikuRAG from haiku.rag.config import load_yaml_config from haiku.rag.config.models import AppConfig -from haiku.rag.tools.context import ToolContextCache +from haiku.rag.skills.rag import create_skill +from haiku.skills.agent import SkillToolset load_dotenv(find_dotenv(usecwd=True)) -# Cache ToolContext instances by thread_id across requests -context_cache = ToolContextCache() - # Configure logfire (only sends data if LOGFIRE_TOKEN is present) try: import logfire @@ -69,34 +62,24 @@ def get_client() -> HaikuRAG: return _client -# Toolkit and agent are created once at module level -chat_toolkit = build_chat_toolkit(Config) -agent = create_chat_agent(Config, toolkit=chat_toolkit) +# Create skill, toolset, and agent +skill = create_skill(db_path=db_path, config=Config) +toolset = SkillToolset(skills=[skill]) +agent = Agent( + os.getenv("HAIKU_CHAT_MODEL", "openai:gpt-4o"), + instructions=toolset.system_prompt, + toolsets=[toolset], +) async def stream_chat(request: Request) -> Response: - """Chat streaming endpoint with AG-UI protocol. - - Uses ToolContextCache to maintain state across requests for the same thread. - AGUIAdapter restores client-sent state via ChatDeps.state setter. - """ + """Chat streaming endpoint with AG-UI protocol.""" body = await request.body() accept = request.headers.get("accept", SSE_CONTENT_TYPE) run_input = AGUIAdapter.build_run_input(body) - thread_id = getattr(run_input, "thread_id", None) or "default" - context, is_new = context_cache.get_or_create(thread_id) - if is_new: - chat_toolkit.prepare(context, state_key=AGUI_STATE_KEY) - - deps = ChatDeps( - config=Config, - client=get_client(), - tool_context=context, - ) - adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept) - event_stream = adapter.run_stream(deps=deps) + event_stream = adapter.run_stream() sse_event_stream = adapter.encode_stream(event_stream) return StreamingResponse( diff --git a/haiku_rag_slim/haiku/rag/chat/__init__.py b/haiku_rag_slim/haiku/rag/chat/__init__.py index 31ee094c..351615e3 100644 --- a/haiku_rag_slim/haiku/rag/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/chat/__init__.py @@ -6,7 +6,7 @@ def run_chat( db_path: Path | None = None, read_only: bool = False, before: datetime | None = None, - initial_context: str | None = None, + model: str | None = None, ) -> None: """Run the chat TUI. @@ -14,7 +14,7 @@ def run_chat( 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. - initial_context: Initial background context to provide to the conversation. + model: Model to use for the chat agent. """ try: from haiku.rag.chat.app import ChatApp @@ -24,15 +24,19 @@ def run_chat( ) from e from haiku.rag.config import get_config + from haiku.rag.skills.rag import create_skill config = get_config() if db_path is None: db_path = config.storage.data_dir / "haiku.rag.lancedb" + skill = create_skill(db_path=db_path, config=config) + app = ChatApp( db_path, + skill=skill, read_only=read_only, before=before, - initial_context=initial_context, + model=model, ) app.run() diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 47552bc8..dd789b74 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -1,30 +1,15 @@ # pyright: reportPossiblyUnboundVariable=false import asyncio import uuid -from collections.abc import AsyncIterable, Iterable +from collections.abc import Iterable from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Any -from pydantic_ai import ( - Agent, - AgentStreamEvent, - FunctionToolCallEvent, - FunctionToolResultEvent, - RunContext, -) -from pydantic_ai.messages import ModelMessage - -from haiku.rag.agents.chat.agent import ( - ChatDeps, - build_chat_toolkit, - create_chat_agent, - trigger_background_summarization, -) -from haiku.rag.agents.chat.state import AGUI_STATE_KEY from haiku.rag.client import HaikuRAG from haiku.rag.config import get_config -from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState +from haiku.skills.agent import SkillToolset +from haiku.skills.models import Skill if TYPE_CHECKING: from textual.app import ComposeResult @@ -39,6 +24,20 @@ except ImportError: # pragma: no cover try: import textual_image.widget # noqa: F401 - import early for renderer detection + from ag_ui.core import ( + AssistantMessage, + BaseEvent, + EventType, + RunAgentInput, + StateDeltaEvent, + TextMessageContentEvent, + ToolCallEndEvent, + ToolCallStartEvent, + UserMessage, + ) + from jsonpatch import JsonPatch + from pydantic_ai import Agent + from pydantic_ai.ag_ui import AGUIAdapter from textual.app import App, SystemCommand from textual.binding import Binding from textual.widgets import Footer, Header, Input @@ -53,6 +52,9 @@ except ImportError: # pragma: no cover SystemCommand = object # type: ignore +RAG_STATE_NAMESPACE = "rag" + + class ChatApp(App): """Textual TUI for conversational RAG.""" @@ -86,23 +88,25 @@ class ChatApp(App): def __init__( self, db_path: Path, + skill: Skill, read_only: bool = False, before: datetime | None = None, - initial_context: str | None = None, + model: str | None = None, ) -> None: super().__init__() self.db_path = db_path + self._skill = skill self.read_only = read_only self.before = before - self._initial_context = initial_context - self._context_locked = False + self._model = model or "openai:gpt-4o" self.client: HaikuRAG | None = None self.config = get_config() - self.agent: Agent[ChatDeps, str] | None = None + self._toolset: SkillToolset | None = None + self._agent: Agent[None, str] | None = None + self._messages: list[Any] = [] + self._state: dict[str, Any] = {} self._is_processing = False - self._tool_call_widgets: dict[str, Any] = {} self._current_worker: Worker[None] | None = None - self._message_history: list[ModelMessage] = [] self._document_filter: list[str] = [] def compose(self) -> "ComposeResult": @@ -136,9 +140,9 @@ class ChatApp(App): self.action_show_info, ) yield SystemCommand( - "Memory", - "View/edit context (editable before first message)", - self.action_show_context, + "View state", + "Show the current session state", + self.action_view_state, ) async def on_mount(self) -> None: @@ -151,17 +155,14 @@ class ChatApp(App): ) await self.client.__aenter__() - # Create toolkit, context, and agent - self.toolkit = build_chat_toolkit(self.config) - self.tool_context = self.toolkit.create_context(state_key=AGUI_STATE_KEY) - self.agent = create_chat_agent(self.config, toolkit=self.toolkit) + self._toolset = SkillToolset(skills=[self._skill]) + self._agent = Agent( + self._model, + instructions=self._toolset.system_prompt, + toolsets=[self._toolset], + ) + self._state = self._toolset.build_state_snapshot() - # Sync document filter to tool context - session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) - if session_state is not None: - session_state.document_filter = self._document_filter - - # Focus the input field self.query_one(Input).focus() async def on_unmount(self) -> None: @@ -169,60 +170,25 @@ class ChatApp(App): 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) - - 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 - - # Lock context after first message - self._context_locked = True - - # 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() - session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) - if session_state: - session_state.citations.clear() + self._messages.append( + UserMessage( + id=str(uuid.uuid4()), + role="user", + content=user_message, + ) + ) - # 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( @@ -231,64 +197,75 @@ class ChatApp(App): async def _run_agent(self, user_message: str) -> None: """Run the agent in a background worker.""" - if not self.client or not self.agent: + if not self._agent or not self._toolset: return chat_history = self.query_one(ChatHistory) - - # Show thinking indicator await chat_history.show_thinking() + run_input = RunAgentInput( + thread_id="tui", + run_id=str(uuid.uuid4()), + messages=self._messages, + state=self._state, + tools=[], + context=[], + forwarded_props={}, + ) + + adapter = AGUIAdapter(agent=self._agent, run_input=run_input) + + message = None + accumulated_text = "" + try: - # Promote initial_context to QA session context on first run - from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState - - qa_session_state = self.tool_context.get( - QA_SESSION_NAMESPACE, QASessionState - ) - if qa_session_state is not None: - if qa_session_state.session_context is None and self._initial_context: - from haiku.rag.tools.session import SessionContext - - qa_session_state.session_context = SessionContext( - summary=self._initial_context + async for event in adapter.run_stream(): + if not isinstance(event, BaseEvent): + continue + if event.type == EventType.TEXT_MESSAGE_START: + chat_history.hide_thinking() + message = await chat_history.add_message("assistant") + accumulated_text = "" + elif event.type == EventType.TEXT_MESSAGE_CONTENT: + assert isinstance(event, TextMessageContentEvent) + accumulated_text += event.delta + if message: + message.update_content(accumulated_text) + chat_history.scroll_end(animate=False) + elif event.type == EventType.TEXT_MESSAGE_END: + self._messages.append( + AssistantMessage( + id=str(uuid.uuid4()), + role="assistant", + content=accumulated_text, + ) ) - - deps = ChatDeps( - config=self.config, - client=self.client, - tool_context=self.tool_context, - ) - - 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 from ToolContext - session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) - if session_state and session_state.citations: - await chat_history.add_citations(session_state.citations) - - # Trigger background summarization - trigger_background_summarization(deps) + # Show citations from RAG state + await self._show_citations(chat_history) + elif event.type == EventType.TOOL_CALL_START: + assert isinstance(event, ToolCallStartEvent) + chat_history.hide_thinking() + await chat_history.add_tool_call( + event.tool_call_id, event.tool_call_name + ) + await chat_history.show_thinking("Executing tasks...") + elif event.type == EventType.TOOL_CALL_END: + assert isinstance(event, ToolCallEndEvent) + chat_history.mark_tool_complete(event.tool_call_id) + elif event.type == EventType.STATE_DELTA: + assert isinstance(event, StateDeltaEvent) + patch = JsonPatch(event.delta) + self._state = patch.apply(self._state) + self._toolset.restore_state_snapshot(self._state) + elif event.type == EventType.STATE_SNAPSHOT: + self._state = getattr(event, "snapshot", self._state) + self._toolset.restore_state_snapshot(self._state) + elif event.type == EventType.RUN_FINISHED: + chat_history.hide_thinking() + elif event.type == EventType.RUN_ERROR: + chat_history.hide_thinking() + error_msg = getattr(event, "message", "Unknown error") + await chat_history.add_message("assistant", f"Error: {error_msg}") except asyncio.CancelledError: chat_history.hide_thinking() @@ -303,20 +280,26 @@ class ChatApp(App): chat_input.disabled = False chat_input.focus() + async def _show_citations(self, chat_history: "ChatHistory") -> None: + """Show citations from the RAG state after an agent response.""" + if not self._toolset: + return + rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE) + if rag_state is None: + return + citations = getattr(rag_state, "citations", []) + if citations: + # Show only new citations (since last response) + await chat_history.add_citations(citations) + async def action_clear_chat(self) -> None: """Clear the chat history and reset session.""" - from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState - chat_history = self.query_one(ChatHistory) await chat_history.clear_messages() - self._message_history.clear() - self._context_locked = False - # Re-register fresh states in ToolContext - self.tool_context.register( - SESSION_NAMESPACE, - SessionState(document_filter=self._document_filter), - ) - self.tool_context.register(QA_SESSION_NAMESPACE, QASessionState()) + self._messages.clear() + # Reset state + if self._toolset: + self._state = self._toolset.build_state_snapshot() def action_focus_input(self) -> None: """Focus the input field, or cancel if processing.""" @@ -340,7 +323,6 @@ class ChatApp(App): if not self.client: return - # Get citation from selected widget directly chat_history = self.query_one(ChatHistory) selected_widgets = list(chat_history.query(CitationWidget).filter(".selected")) if not selected_widgets: @@ -364,38 +346,19 @@ class ChatApp(App): await self.push_screen(InfoModal(self.client, self.db_path)) - async def action_show_context(self) -> None: - """Show context modal (edit initial context or view session context).""" - from haiku.rag.chat.widgets.context_modal import ContextModal - from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState + def action_view_state(self) -> None: + """Show the current session state.""" + from haiku.skills.chat.app import StateScreen - session_context = None - qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) - if qa_session_state and qa_session_state.session_context is not None: - session_context = qa_session_state.session_context - - await self.push_screen( - ContextModal( - initial_context=self._initial_context, - session_context=session_context, - is_locked=self._context_locked, - ) - ) - - def on_context_modal_context_updated(self, event: Any) -> None: - """Handle context updates from modal.""" - if not self._context_locked: - self._initial_context = event.context or None + self.push_screen(StateScreen(self._state)) 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 widget that was focused event.widget.add_class("selected") async def action_show_filter(self) -> None: @@ -415,6 +378,3 @@ class ChatApp(App): def on_document_filter_modal_filter_changed(self, event: Any) -> None: """Handle document filter changes from modal.""" self._document_filter = event.selected - session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) - if session_state: - session_state.document_filter = self._document_filter diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py b/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py index 328e1087..90eb7c85 100644 --- a/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py +++ b/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from textual.containers import Horizontal, VerticalScroll from textual.message import Message @@ -32,52 +32,48 @@ class ChatMessage(Static): class ToolCallWidget(Static): - """Styled inline display of a tool call.""" + """Displays a single tool call with status indicator.""" - TOOL_LABELS = { - "search": "Searching", - "ask": "Asking", - "get_document": "Fetching", - } - - def __init__(self, tool_name: str, args: dict | None = None, **kwargs) -> None: + def __init__( + self, + tool_call_id: str, + tool_name: str, + args: dict[str, Any] | None = None, + **kwargs, + ) -> None: super().__init__(**kwargs) + self.tool_call_id = tool_call_id self.tool_name = tool_name self.args = args or {} - self._complete = False + self._completed = 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: + if self._completed: yield Static("✓", classes="tool-status") else: yield LoadingIndicator(classes="tool-spinner") - yield Static(label, classes="tool-badge") - yield Static(desc, classes="tool-desc") + yield Static(self.tool_name, classes="tool-badge") + desc = self._build_description() + if desc: + yield Static(desc, classes="tool-desc") - def mark_complete(self) -> None: - self._complete = True + def _build_description(self) -> str: + if self.tool_name == "search": + query = self.args.get("query", "...") + return f'"{query}"' + elif self.tool_name == "ask": + question = self.args.get("question", "...") + return f'"{question}"' + elif self.tool_name == "get_document": + query = self.args.get("query", "...") + return f'"{query}"' + elif self.args: + return str(self.args) + return "" + + def mark_completed(self) -> None: + self._completed = True self.refresh(recompose=True) @@ -102,7 +98,6 @@ class CitationWidget(Collapsible): pages += "..." title += f" (p.{pages})" - # Build content widgets content = citation.content if len(content) > 500: content = content[:500] + "..." @@ -132,10 +127,22 @@ class CitationWidget(Collapsible): class ThinkingWidget(Static): """Thinking indicator shown while agent is processing.""" + def __init__(self, text: str = "Thinking...", **kwargs) -> None: + super().__init__(**kwargs) + self._text = text + def compose(self) -> "ComposeResult": with Horizontal(classes="thinking-row"): yield LoadingIndicator(classes="thinking-spinner") - yield Static("Thinking...", classes="thinking-text") + yield Static(self._text, classes="thinking-text", id="thinking-label") + + def update_text(self, text: str) -> None: + self._text = text + try: + label = self.query_one("#thinking-label", Static) + label.update(text) + except Exception: + pass class SourcesHeader(Static): @@ -303,6 +310,7 @@ class ChatHistory(VerticalScroll): def __init__(self, **kwargs) -> None: super().__init__(**kwargs) self.messages: list[tuple[str, str]] = [] + self._tool_widgets: dict[str, ToolCallWidget] = {} async def add_message(self, role: str, content: str = "") -> ChatMessage: """Add a message to the chat history.""" @@ -313,18 +321,24 @@ class ChatHistory(VerticalScroll): return message_widget async def add_tool_call( - self, tool_name: str, args: dict | None = None + self, + tool_call_id: str, + tool_name: str, + args: dict[str, Any] | None = None, ) -> ToolCallWidget: """Add an inline tool call indicator.""" - widget = ToolCallWidget(tool_name, args) + widget = ToolCallWidget(tool_call_id, tool_name, args) + self._tool_widgets[tool_call_id] = widget 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") + def mark_tool_complete(self, tool_call_id: str) -> None: + """Mark a tool call as complete by its ID.""" + widget = self._tool_widgets.get(tool_call_id) + if widget: + widget.mark_completed() + widget.add_class("complete") async def add_citations(self, citations: list[Citation]) -> None: """Add citations inline after a response.""" @@ -336,9 +350,12 @@ class ChatHistory(VerticalScroll): await self.mount(widget) self.scroll_end(animate=False) - async def show_thinking(self) -> None: + async def show_thinking(self, text: str = "Thinking...") -> None: """Show the thinking indicator.""" - await self.mount(ThinkingWidget(id="thinking")) + try: + self.query_one("#thinking", ThinkingWidget).update_text(text) + except Exception: + await self.mount(ThinkingWidget(text, id="thinking")) self.scroll_end(animate=False) def hide_thinking(self) -> None: @@ -351,4 +368,5 @@ class ChatHistory(VerticalScroll): async def clear_messages(self) -> None: """Clear all messages from the chat history.""" self.messages.clear() + self._tool_widgets.clear() await self.remove_children() diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py b/haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py index 92a17316..6298bfb5 100644 --- a/haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py +++ b/haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py @@ -1,22 +1,12 @@ -from typing import TYPE_CHECKING - from textual.app import ComposeResult from textual.binding import Binding from textual.containers import Horizontal, Vertical, VerticalScroll -from textual.message import Message from textual.screen import ModalScreen -from textual.widgets import Button, Markdown, Static, TextArea - -if TYPE_CHECKING: - from haiku.rag.tools.session import SessionContext +from textual.widgets import Button, Markdown, Static class ContextModal(ModalScreen): # pragma: no cover - """Modal screen for viewing/editing context. - - Before first message (not locked, no session context): Edit initial context - After first message (locked or has session context): View session context - """ + """Modal screen for viewing session Q&A history.""" BINDINGS = [ Binding("escape", "cancel", "Close", show=False), @@ -49,12 +39,6 @@ class ContextModal(ModalScreen): # pragma: no cover color: $text-muted; } - #context-editor { - height: 12; - min-height: 8; - max-height: 16; - } - #context-content { height: 1fr; max-height: 16; @@ -73,81 +57,39 @@ class ContextModal(ModalScreen): # pragma: no cover } """ - class ContextUpdated(Message): - """Emitted when the context is saved.""" - - def __init__(self, context: str) -> None: - super().__init__() - self.context = context - - def __init__( - self, - initial_context: str | None = None, - session_context: "SessionContext | None" = None, - is_locked: bool = False, - ) -> None: + def __init__(self, qa_history: list | None = None) -> None: super().__init__() - self._initial_context = initial_context - self._session_context = session_context - self._is_locked = is_locked - - @property - def _is_edit_mode(self) -> bool: - """Edit mode when not locked and no session context yet.""" - has_session_context = self._session_context and self._session_context.summary - return not self._is_locked and not has_session_context + self._qa_history = qa_history or [] def compose(self) -> ComposeResult: with Vertical(id="context-container"): - if self._is_edit_mode: - yield Static("[bold]Initial Context[/bold]", id="context-header") - yield Static( - "Set background context to guide the conversation. " - "This will be locked after you send your first message.", - id="context-description", - ) - initial_value = self._initial_context or "" - yield TextArea(initial_value, id="context-editor") - with Horizontal(id="button-row"): - yield Button("Cancel", id="cancel-btn", variant="default") - yield Button("Save", id="save-btn", variant="primary") - else: - yield Static("[bold]Session Context[/bold]", id="context-header") - yield Static( - "What the assistant has learned from your conversation.", - id="context-description", - ) - with VerticalScroll(id="context-content"): - yield Markdown(self._get_session_content()) - with Horizontal(id="button-row"): - yield Button("Close", id="cancel-btn", variant="primary") + yield Static("[bold]Session Context[/bold]", id="context-header") + yield Static( + "Questions and answers from this session.", + id="context-description", + ) + with VerticalScroll(id="context-content"): + yield Markdown(self._get_content()) + with Horizontal(id="button-row"): + yield Button("Close", id="cancel-btn", variant="primary") - def _get_session_content(self) -> str: - if not self._session_context: - return "*No session context yet. Ask a question first.*" + def _get_content(self) -> str: + if not self._qa_history: + return "*No questions asked yet.*" - ctx = self._session_context - updated = ( - ctx.last_updated.strftime("%Y-%m-%d %H:%M:%S") - if ctx.last_updated - else "unknown" - ) + parts = [] + for entry in self._qa_history: + q = getattr(entry, "question", str(entry)) + a = getattr(entry, "answer", "") + parts.append(f"**Q:** {q}\n\n**A:** {a}") - return f"**Last updated:** {updated}\n\n---\n\n{ctx.summary}" + return "\n\n---\n\n".join(parts) def on_button_pressed(self, event: Button.Pressed) -> None: """Handle button presses.""" if event.button.id == "cancel-btn": self.action_cancel() - elif event.button.id == "save-btn": - self.action_save() def action_cancel(self) -> None: - """Cancel and close without saving.""" - self.app.pop_screen() - - def action_save(self) -> None: - """Save context and close.""" - editor = self.query_one("#context-editor", TextArea) - self.post_message(self.ContextUpdated(editor.text)) + """Cancel and close.""" self.app.pop_screen() diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index f4053122..f8aa3a6a 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -620,10 +620,10 @@ def chat( # pragma: no cover "--db", help="Path to the LanceDB database file", ), - initial_context: str | None = typer.Option( + model: str | None = typer.Option( None, - "--initial-context", - help="Initial background context to provide to the conversation", + "--model", + help="Model to use for the chat agent (e.g. openai:gpt-4o)", ), ): """Launch the chat TUI for conversational RAG.""" @@ -635,7 +635,7 @@ def chat( # pragma: no cover db_path, read_only=_read_only, before=_before, - initial_context=initial_context, + model=model, ) diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 48f78d44..1c41e9b1 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -5,7 +5,6 @@ import pytest from typer.testing import CliRunner from haiku.rag.cli import _cli as cli -from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState runner = CliRunner() @@ -21,19 +20,46 @@ def test_chat_command(): 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 - +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) + 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() + + skill = MagicMock() + skill.state_type = None + skill.state_namespace = None + skill.tools = [] + skill.toolsets = [] + skill.resources = [] + skill.metadata = MagicMock() + skill.metadata.name = "rag" + skill.metadata.description = "RAG skill" + + return ChatApp( + db_path=db_path, + skill=skill, + 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, Input.""" + from haiku.rag.chat.widgets.chat_history import ChatHistory + + app, mock_client = _make_app(temp_db_path) 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 @@ -47,144 +73,64 @@ async def test_chat_app_has_required_widgets(temp_db_path: Path): @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) + app, mock_client = _make_app(temp_db_path) 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) + app, mock_client = _make_app(temp_db_path) 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) + app, mock_client = _make_app(temp_db_path) 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"}) + tool_widget = await chat_history.add_tool_call( + "tool-1", "search", {"query": "test"} + ) assert isinstance(tool_widget, ToolCallWidget) - assert tool_widget._complete is False + assert tool_widget._completed is False - # Mark it complete - chat_history.mark_tool_complete(tool_widget) - assert tool_widget._complete is True + 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.agents.research.models import Citation - 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) + app, mock_client = _make_app(temp_db_path) 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) @@ -213,7 +159,6 @@ async def test_chat_history_can_add_citations(temp_db_path: Path): await chat_history.add_citations(test_citations) - # Verify citation widgets were added citation_widgets = chat_history.query(CitationWidget) assert len(list(citation_widgets)) == 2 @@ -221,25 +166,18 @@ async def test_chat_history_can_add_citations(temp_db_path: Path): @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) + app, mock_client = _make_app(temp_db_path) 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) @@ -247,67 +185,38 @@ async def test_chat_history_thinking_indicator(temp_db_path: Path): @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 +async def test_clear_chat_resets_state(temp_db_path: Path): + """Test that clearing chat resets state and messages.""" 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) + app, mock_client = _make_app(temp_db_path) 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 - # Clear chat via action (available through command palette) await app.action_clear_chat() await pilot.pause() - # Verify messages cleared assert len(chat_history.messages) == 0 - # Verify ToolContext states are reset - from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState - - session_state = app.tool_context.get(SESSION_NAMESPACE, SessionState) - assert session_state is not None - assert session_state.citations == [] - assert session_state.citation_registry == {} - - qa_session_state = app.tool_context.get( - QA_SESSION_NAMESPACE, QASessionState - ) - assert qa_session_state is not None - assert qa_session_state.qa_history == [] - assert qa_session_state.session_context is None - @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.research.models import Citation - 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) + app, mock_client = _make_app(temp_db_path) 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 = Citation( index=1, document_id="doc1", @@ -319,20 +228,16 @@ async def test_citation_expand_collapse_with_enter(temp_db_path: Path): ) 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