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

Rewrite chat TUI and app backend with haiku.skills

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-19 17:01:17 +02:00
parent 8176103eb0
commit 643ed20d6f
No known key found for this signature in database
7 changed files with 289 additions and 477 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
app/backend/main.py
View file

@ -3,8 +3,9 @@ import os
from pathlib import Path from pathlib import Path
from dotenv import find_dotenv, load_dotenv 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 import SSE_CONTENT_TYPE
from pydantic_ai.ui.ag_ui import AGUIAdapter
from starlette.applications import Starlette from starlette.applications import Starlette
from starlette.middleware import Middleware from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware from starlette.middleware.cors import CORSMiddleware
@ -12,22 +13,14 @@ from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route 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.client import HaikuRAG
from haiku.rag.config import load_yaml_config from haiku.rag.config import load_yaml_config
from haiku.rag.config.models import AppConfig 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)) 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) # Configure logfire (only sends data if LOGFIRE_TOKEN is present)
try: try:
import logfire import logfire
@ -69,34 +62,24 @@ def get_client() -> HaikuRAG:
return _client return _client
# Toolkit and agent are created once at module level # Create skill, toolset, and agent
chat_toolkit = build_chat_toolkit(Config) skill = create_skill(db_path=db_path, config=Config)
agent = create_chat_agent(Config, toolkit=chat_toolkit) 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: async def stream_chat(request: Request) -> Response:
"""Chat streaming endpoint with AG-UI protocol. """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.
"""
body = await request.body() body = await request.body()
accept = request.headers.get("accept", SSE_CONTENT_TYPE) accept = request.headers.get("accept", SSE_CONTENT_TYPE)
run_input = AGUIAdapter.build_run_input(body) 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) 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) sse_event_stream = adapter.encode_stream(event_stream)
return StreamingResponse( return StreamingResponse(

10
haiku_rag_slim/haiku/rag/chat/__init__.py
View file

@ -6,7 +6,7 @@ def run_chat(
db_path: Path | None = None, db_path: Path | None = None,
read_only: bool = False, read_only: bool = False,
before: datetime | None = None, before: datetime | None = None,
initial_context: str | None = None, model: str | None = None,
) -> None: ) -> None:
"""Run the chat TUI. """Run the chat TUI.
@ -14,7 +14,7 @@ def run_chat(
db_path: Path to the LanceDB database. If None, uses default from config. db_path: Path to the LanceDB database. If None, uses default from config.
read_only: Whether to open the database in read-only mode. read_only: Whether to open the database in read-only mode.
before: Query database as it existed before this datetime. 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: try:
from haiku.rag.chat.app import ChatApp from haiku.rag.chat.app import ChatApp
@ -24,15 +24,19 @@ def run_chat(
) from e ) from e
from haiku.rag.config import get_config from haiku.rag.config import get_config
from haiku.rag.skills.rag import create_skill
config = get_config() config = get_config()
if db_path is None: if db_path is None:
db_path = config.storage.data_dir / "haiku.rag.lancedb" db_path = config.storage.data_dir / "haiku.rag.lancedb"
skill = create_skill(db_path=db_path, config=config)
app = ChatApp( app = ChatApp(
db_path, db_path,
skill=skill,
read_only=read_only, read_only=read_only,
before=before, before=before,
initial_context=initial_context, model=model,
) )
app.run() app.run()

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

@ -1,30 +1,15 @@
# pyright: reportPossiblyUnboundVariable=false # pyright: reportPossiblyUnboundVariable=false
import asyncio import asyncio
import uuid import uuid
from collections.abc import AsyncIterable, Iterable from collections.abc import Iterable
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any 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.client import HaikuRAG
from haiku.rag.config import get_config 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: if TYPE_CHECKING:
from textual.app import ComposeResult from textual.app import ComposeResult
@ -39,6 +24,20 @@ except ImportError: # pragma: no cover
try: try:
import textual_image.widget # noqa: F401 - import early for renderer detection 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.app import App, SystemCommand
from textual.binding import Binding from textual.binding import Binding
from textual.widgets import Footer, Header, Input from textual.widgets import Footer, Header, Input
@ -53,6 +52,9 @@ except ImportError: # pragma: no cover
SystemCommand = object # type: ignore SystemCommand = object # type: ignore
RAG_STATE_NAMESPACE = "rag"
class ChatApp(App): class ChatApp(App):
"""Textual TUI for conversational RAG.""" """Textual TUI for conversational RAG."""
@ -86,23 +88,25 @@ class ChatApp(App):
def __init__( def __init__(
self, self,
db_path: Path, db_path: Path,
skill: Skill,
read_only: bool = False, read_only: bool = False,
before: datetime | None = None, before: datetime | None = None,
initial_context: str | None = None, model: str | None = None,
) -> None: ) -> None:
super().__init__() super().__init__()
self.db_path = db_path self.db_path = db_path
self._skill = skill
self.read_only = read_only self.read_only = read_only
self.before = before self.before = before
self._initial_context = initial_context self._model = model or "openai:gpt-4o"
self._context_locked = False
self.client: HaikuRAG | None = None self.client: HaikuRAG | None = None
self.config = get_config() 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._is_processing = False
self._tool_call_widgets: dict[str, Any] = {}
self._current_worker: Worker[None] | None = None self._current_worker: Worker[None] | None = None
self._message_history: list[ModelMessage] = []
self._document_filter: list[str] = [] self._document_filter: list[str] = []
def compose(self) -> "ComposeResult": def compose(self) -> "ComposeResult":
@ -136,9 +140,9 @@ class ChatApp(App):
self.action_show_info, self.action_show_info,
) )
yield SystemCommand( yield SystemCommand(
"Memory", "View state",
"View/edit context (editable before first message)", "Show the current session state",
self.action_show_context, self.action_view_state,
) )
async def on_mount(self) -> None: async def on_mount(self) -> None:
@ -151,17 +155,14 @@ class ChatApp(App):
) )
await self.client.__aenter__() await self.client.__aenter__()
# Create toolkit, context, and agent self._toolset = SkillToolset(skills=[self._skill])
self.toolkit = build_chat_toolkit(self.config) self._agent = Agent(
self.tool_context = self.toolkit.create_context(state_key=AGUI_STATE_KEY) self._model,
self.agent = create_chat_agent(self.config, toolkit=self.toolkit) 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() self.query_one(Input).focus()
async def on_unmount(self) -> None: async def on_unmount(self) -> None:
@ -169,60 +170,25 @@ class ChatApp(App):
if self.client: if self.client:
await self.client.__aexit__(None, None, None) 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: async def on_input_submitted(self, event: Input.Submitted) -> None:
"""Handle user input submission.""" """Handle user input submission."""
user_message = event.value.strip() user_message = event.value.strip()
if not user_message or self._is_processing: if not user_message or self._is_processing:
return 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() event.input.clear()
# Add user message to history
chat_history = self.query_one(ChatHistory) chat_history = self.query_one(ChatHistory)
await chat_history.add_message("user", user_message) await chat_history.add_message("user", user_message)
# Clear for new query self._messages.append(
self._tool_call_widgets.clear() UserMessage(
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) id=str(uuid.uuid4()),
if session_state: role="user",
session_state.citations.clear() content=user_message,
)
)
# Run agent in a worker to keep UI responsive
self._is_processing = True self._is_processing = True
self.query_one(Input).disabled = True self.query_one(Input).disabled = True
self._current_worker = self.run_worker( self._current_worker = self.run_worker(
@ -231,64 +197,75 @@ class ChatApp(App):
async def _run_agent(self, user_message: str) -> None: async def _run_agent(self, user_message: str) -> None:
"""Run the agent in a background worker.""" """Run the agent in a background worker."""
if not self.client or not self.agent: if not self._agent or not self._toolset:
return return
chat_history = self.query_one(ChatHistory) chat_history = self.query_one(ChatHistory)
# Show thinking indicator
await chat_history.show_thinking() 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: try:
# Promote initial_context to QA session context on first run async for event in adapter.run_stream():
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState if not isinstance(event, BaseEvent):
continue
qa_session_state = self.tool_context.get( if event.type == EventType.TEXT_MESSAGE_START:
QA_SESSION_NAMESPACE, QASessionState chat_history.hide_thinking()
) message = await chat_history.add_message("assistant")
if qa_session_state is not None: accumulated_text = ""
if qa_session_state.session_context is None and self._initial_context: elif event.type == EventType.TEXT_MESSAGE_CONTENT:
from haiku.rag.tools.session import SessionContext assert isinstance(event, TextMessageContentEvent)
accumulated_text += event.delta
qa_session_state.session_context = SessionContext( if message:
summary=self._initial_context 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,
)
) )
# Show citations from RAG state
deps = ChatDeps( await self._show_citations(chat_history)
config=self.config, elif event.type == EventType.TOOL_CALL_START:
client=self.client, assert isinstance(event, ToolCallStartEvent)
tool_context=self.tool_context, chat_history.hide_thinking()
) await chat_history.add_tool_call(
event.tool_call_id, event.tool_call_name
async with self.agent.run_stream( )
user_message, await chat_history.show_thinking("Executing tasks...")
deps=deps, elif event.type == EventType.TOOL_CALL_END:
message_history=self._message_history, assert isinstance(event, ToolCallEndEvent)
event_stream_handler=self._event_stream_handler, chat_history.mark_tool_complete(event.tool_call_id)
) as stream: elif event.type == EventType.STATE_DELTA:
# Hide thinking when we start getting content assert isinstance(event, StateDeltaEvent)
chat_history.hide_thinking() patch = JsonPatch(event.delta)
self._state = patch.apply(self._state)
# Create assistant message for streaming self._toolset.restore_state_snapshot(self._state)
assistant_msg = await chat_history.add_message("assistant", "") elif event.type == EventType.STATE_SNAPSHOT:
self._state = getattr(event, "snapshot", self._state)
# Stream text updates self._toolset.restore_state_snapshot(self._state)
async for text in stream.stream_text(): elif event.type == EventType.RUN_FINISHED:
assistant_msg.update_content(text) chat_history.hide_thinking()
chat_history.scroll_end(animate=False) elif event.type == EventType.RUN_ERROR:
# Yield to event loop to keep UI responsive chat_history.hide_thinking()
await asyncio.sleep(0) error_msg = getattr(event, "message", "Unknown error")
await chat_history.add_message("assistant", f"Error: {error_msg}")
# 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)
except asyncio.CancelledError: except asyncio.CancelledError:
chat_history.hide_thinking() chat_history.hide_thinking()
@ -303,20 +280,26 @@ class ChatApp(App):
chat_input.disabled = False chat_input.disabled = False
chat_input.focus() 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: async def action_clear_chat(self) -> None:
"""Clear the chat history and reset session.""" """Clear the chat history and reset session."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
chat_history = self.query_one(ChatHistory) chat_history = self.query_one(ChatHistory)
await chat_history.clear_messages() await chat_history.clear_messages()
self._message_history.clear() self._messages.clear()
self._context_locked = False # Reset state
# Re-register fresh states in ToolContext if self._toolset:
self.tool_context.register( self._state = self._toolset.build_state_snapshot()
SESSION_NAMESPACE,
SessionState(document_filter=self._document_filter),
)
self.tool_context.register(QA_SESSION_NAMESPACE, QASessionState())
def action_focus_input(self) -> None: def action_focus_input(self) -> None:
"""Focus the input field, or cancel if processing.""" """Focus the input field, or cancel if processing."""
@ -340,7 +323,6 @@ class ChatApp(App):
if not self.client: if not self.client:
return return
# Get citation from selected widget directly
chat_history = self.query_one(ChatHistory) chat_history = self.query_one(ChatHistory)
selected_widgets = list(chat_history.query(CitationWidget).filter(".selected")) selected_widgets = list(chat_history.query(CitationWidget).filter(".selected"))
if not selected_widgets: if not selected_widgets:
@ -364,38 +346,19 @@ class ChatApp(App):
await self.push_screen(InfoModal(self.client, self.db_path)) await self.push_screen(InfoModal(self.client, self.db_path))
async def action_show_context(self) -> None: def action_view_state(self) -> None:
"""Show context modal (edit initial context or view session context).""" """Show the current session state."""
from haiku.rag.chat.widgets.context_modal import ContextModal from haiku.skills.chat.app import StateScreen
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
session_context = None self.push_screen(StateScreen(self._state))
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
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None: def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
"""Handle citation selection.""" """Handle citation selection."""
chat_history = self.query_one(ChatHistory) chat_history = self.query_one(ChatHistory)
# Remove selected class from all citations
for widget in chat_history.query(CitationWidget): for widget in chat_history.query(CitationWidget):
widget.remove_class("selected") widget.remove_class("selected")
# Add selected class to the widget that was focused
event.widget.add_class("selected") event.widget.add_class("selected")
async def action_show_filter(self) -> None: 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: def on_document_filter_modal_filter_changed(self, event: Any) -> None:
"""Handle document filter changes from modal.""" """Handle document filter changes from modal."""
self._document_filter = event.selected self._document_filter = event.selected
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
if session_state:
session_state.document_filter = self._document_filter

110
haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py
View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING from typing import TYPE_CHECKING, Any
from textual.containers import Horizontal, VerticalScroll from textual.containers import Horizontal, VerticalScroll
from textual.message import Message from textual.message import Message
@ -32,52 +32,48 @@ class ChatMessage(Static):
class ToolCallWidget(Static): class ToolCallWidget(Static):
"""Styled inline display of a tool call.""" """Displays a single tool call with status indicator."""
TOOL_LABELS = { def __init__(
"search": "Searching", self,
"ask": "Asking", tool_call_id: str,
"get_document": "Fetching", tool_name: str,
} args: dict[str, Any] | None = None,
**kwargs,
def __init__(self, tool_name: str, args: dict | None = None, **kwargs) -> None: ) -> None:
super().__init__(**kwargs) super().__init__(**kwargs)
self.tool_call_id = tool_call_id
self.tool_name = tool_name self.tool_name = tool_name
self.args = args or {} self.args = args or {}
self._complete = False self._completed = False
def compose(self) -> "ComposeResult": 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"): with Horizontal(classes="tool-row"):
if self._complete: if self._completed:
yield Static("✓", classes="tool-status") yield Static("✓", classes="tool-status")
else: else:
yield LoadingIndicator(classes="tool-spinner") yield LoadingIndicator(classes="tool-spinner")
yield Static(label, classes="tool-badge") yield Static(self.tool_name, classes="tool-badge")
yield Static(desc, classes="tool-desc") desc = self._build_description()
if desc:
yield Static(desc, classes="tool-desc")
def mark_complete(self) -> None: def _build_description(self) -> str:
self._complete = True 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) self.refresh(recompose=True)
@ -102,7 +98,6 @@ class CitationWidget(Collapsible):
pages += "..." pages += "..."
title += f" (p.{pages})" title += f" (p.{pages})"
# Build content widgets
content = citation.content content = citation.content
if len(content) > 500: if len(content) > 500:
content = content[:500] + "..." content = content[:500] + "..."
@ -132,10 +127,22 @@ class CitationWidget(Collapsible):
class ThinkingWidget(Static): class ThinkingWidget(Static):
"""Thinking indicator shown while agent is processing.""" """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": def compose(self) -> "ComposeResult":
with Horizontal(classes="thinking-row"): with Horizontal(classes="thinking-row"):
yield LoadingIndicator(classes="thinking-spinner") 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): class SourcesHeader(Static):
@ -303,6 +310,7 @@ class ChatHistory(VerticalScroll):
def __init__(self, **kwargs) -> None: def __init__(self, **kwargs) -> None:
super().__init__(**kwargs) super().__init__(**kwargs)
self.messages: list[tuple[str, str]] = [] self.messages: list[tuple[str, str]] = []
self._tool_widgets: dict[str, ToolCallWidget] = {}
async def add_message(self, role: str, content: str = "") -> ChatMessage: async def add_message(self, role: str, content: str = "") -> ChatMessage:
"""Add a message to the chat history.""" """Add a message to the chat history."""
@ -313,18 +321,24 @@ class ChatHistory(VerticalScroll):
return message_widget return message_widget
async def add_tool_call( 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: ) -> ToolCallWidget:
"""Add an inline tool call indicator.""" """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) await self.mount(widget)
self.scroll_end(animate=False) self.scroll_end(animate=False)
return widget return widget
def mark_tool_complete(self, widget: ToolCallWidget) -> None: def mark_tool_complete(self, tool_call_id: str) -> None:
"""Mark a tool call as complete.""" """Mark a tool call as complete by its ID."""
widget.mark_complete() widget = self._tool_widgets.get(tool_call_id)
widget.add_class("complete") if widget:
widget.mark_completed()
widget.add_class("complete")
async def add_citations(self, citations: list[Citation]) -> None: async def add_citations(self, citations: list[Citation]) -> None:
"""Add citations inline after a response.""" """Add citations inline after a response."""
@ -336,9 +350,12 @@ class ChatHistory(VerticalScroll):
await self.mount(widget) await self.mount(widget)
self.scroll_end(animate=False) self.scroll_end(animate=False)
async def show_thinking(self) -> None: async def show_thinking(self, text: str = "Thinking...") -> None:
"""Show the thinking indicator.""" """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) self.scroll_end(animate=False)
def hide_thinking(self) -> None: def hide_thinking(self) -> None:
@ -351,4 +368,5 @@ class ChatHistory(VerticalScroll):
async def clear_messages(self) -> None: async def clear_messages(self) -> None:
"""Clear all messages from the chat history.""" """Clear all messages from the chat history."""
self.messages.clear() self.messages.clear()
self._tool_widgets.clear()
await self.remove_children() await self.remove_children()

104
haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py
View file

@ -1,22 +1,12 @@
from typing import TYPE_CHECKING
from textual.app import ComposeResult from textual.app import ComposeResult
from textual.binding import Binding from textual.binding import Binding
from textual.containers import Horizontal, Vertical, VerticalScroll from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.message import Message
from textual.screen import ModalScreen from textual.screen import ModalScreen
from textual.widgets import Button, Markdown, Static, TextArea from textual.widgets import Button, Markdown, Static
if TYPE_CHECKING:
from haiku.rag.tools.session import SessionContext
class ContextModal(ModalScreen): # pragma: no cover class ContextModal(ModalScreen): # pragma: no cover
"""Modal screen for viewing/editing context. """Modal screen for viewing session Q&A history."""
Before first message (not locked, no session context): Edit initial context
After first message (locked or has session context): View session context
"""
BINDINGS = [ BINDINGS = [
Binding("escape", "cancel", "Close", show=False), Binding("escape", "cancel", "Close", show=False),
@ -49,12 +39,6 @@ class ContextModal(ModalScreen): # pragma: no cover
color: $text-muted; color: $text-muted;
} }
#context-editor {
height: 12;
min-height: 8;
max-height: 16;
}
#context-content { #context-content {
height: 1fr; height: 1fr;
max-height: 16; max-height: 16;
@ -73,81 +57,39 @@ class ContextModal(ModalScreen): # pragma: no cover
} }
""" """
class ContextUpdated(Message): def __init__(self, qa_history: list | None = None) -> None:
"""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:
super().__init__() super().__init__()
self._initial_context = initial_context self._qa_history = qa_history or []
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
def compose(self) -> ComposeResult: def compose(self) -> ComposeResult:
with Vertical(id="context-container"): with Vertical(id="context-container"):
if self._is_edit_mode: yield Static("[bold]Session Context[/bold]", id="context-header")
yield Static("[bold]Initial Context[/bold]", id="context-header") yield Static(
yield Static( "Questions and answers from this session.",
"Set background context to guide the conversation. " id="context-description",
"This will be locked after you send your first message.", )
id="context-description", with VerticalScroll(id="context-content"):
) yield Markdown(self._get_content())
initial_value = self._initial_context or "" with Horizontal(id="button-row"):
yield TextArea(initial_value, id="context-editor") yield Button("Close", id="cancel-btn", variant="primary")
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")
def _get_session_content(self) -> str: def _get_content(self) -> str:
if not self._session_context: if not self._qa_history:
return "*No session context yet. Ask a question first.*" return "*No questions asked yet.*"
ctx = self._session_context parts = []
updated = ( for entry in self._qa_history:
ctx.last_updated.strftime("%Y-%m-%d %H:%M:%S") q = getattr(entry, "question", str(entry))
if ctx.last_updated a = getattr(entry, "answer", "")
else "unknown" 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: def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button presses.""" """Handle button presses."""
if event.button.id == "cancel-btn": if event.button.id == "cancel-btn":
self.action_cancel() self.action_cancel()
elif event.button.id == "save-btn":
self.action_save()
def action_cancel(self) -> None: def action_cancel(self) -> None:
"""Cancel and close without saving.""" """Cancel and close."""
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))
self.app.pop_screen() self.app.pop_screen()

8
haiku_rag_slim/haiku/rag/cli.py
View file

@ -620,10 +620,10 @@ def chat( # pragma: no cover
"--db", "--db",
help="Path to the LanceDB database file", help="Path to the LanceDB database file",
), ),
initial_context: str | None = typer.Option( model: str | None = typer.Option(
None, None,
"--initial-context", "--model",
help="Initial background context to provide to the conversation", help="Model to use for the chat agent (e.g. openai:gpt-4o)",
), ),
): ):
"""Launch the chat TUI for conversational RAG.""" """Launch the chat TUI for conversational RAG."""
@ -635,7 +635,7 @@ def chat( # pragma: no cover
db_path, db_path,
read_only=_read_only, read_only=_read_only,
before=_before, before=_before,
initial_context=initial_context, model=model,
) )

195
tests/chat/test_chat_app.py
View file

@ -5,7 +5,6 @@ import pytest
from typer.testing import CliRunner from typer.testing import CliRunner
from haiku.rag.cli import _cli as cli from haiku.rag.cli import _cli as cli
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
runner = CliRunner() runner = CliRunner()
@ -21,19 +20,46 @@ def test_chat_command():
mock_chat.assert_called_once() mock_chat.assert_called_once()
@pytest.mark.asyncio def _make_mock_client():
async def test_chat_app_has_required_widgets(temp_db_path: Path): """Create a mock HaikuRAG client."""
"""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 = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None) 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): with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
app = ChatApp(temp_db_path, read_only=True)
async with app.run_test(): async with app.run_test():
chat_history = app.query_one(ChatHistory) chat_history = app.query_one(ChatHistory)
assert chat_history is not None 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 @pytest.mark.asyncio
async def test_chat_app_quit_binding(temp_db_path: Path): async def test_chat_app_quit_binding(temp_db_path: Path):
"""Test that pressing ctrl+q quits the app.""" """Test that pressing ctrl+q quits the app."""
from haiku.rag.chat.app import ChatApp app, mock_client = _make_app(temp_db_path)
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): 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: async with app.run_test() as pilot:
# App should be running
assert app.is_running assert app.is_running
# Press ctrl+q to quit
await pilot.press("ctrl+q") await pilot.press("ctrl+q")
# App should have exited
assert not app.is_running assert not app.is_running
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_chat_history_can_add_message(temp_db_path: Path): async def test_chat_history_can_add_message(temp_db_path: Path):
"""Test that ChatHistory can display messages.""" """Test that ChatHistory can display messages."""
from haiku.rag.chat.app import ChatApp
from haiku.rag.chat.widgets.chat_history import ChatHistory from haiku.rag.chat.widgets.chat_history import ChatHistory
mock_client = AsyncMock() app, mock_client = _make_app(temp_db_path)
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): with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
app = ChatApp(temp_db_path, read_only=True)
async with app.run_test(): async with app.run_test():
chat_history = app.query_one(ChatHistory) chat_history = app.query_one(ChatHistory)
# Add a user message
await chat_history.add_message("user", "Hello, how are you?") await chat_history.add_message("user", "Hello, how are you?")
assert len(chat_history.messages) == 1 assert len(chat_history.messages) == 1
assert chat_history.messages[0] == ("user", "Hello, how are you?") 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!") await chat_history.add_message("assistant", "I'm doing well, thank you!")
assert len(chat_history.messages) == 2 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 @pytest.mark.asyncio
async def test_chat_history_can_add_tool_calls(temp_db_path: Path): async def test_chat_history_can_add_tool_calls(temp_db_path: Path):
"""Test that ChatHistory can display inline tool calls.""" """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 from haiku.rag.chat.widgets.chat_history import ChatHistory, ToolCallWidget
mock_client = AsyncMock() app, mock_client = _make_app(temp_db_path)
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): with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
app = ChatApp(temp_db_path, read_only=True)
async with app.run_test(): async with app.run_test():
chat_history = app.query_one(ChatHistory) chat_history = app.query_one(ChatHistory)
# Add a tool call tool_widget = await chat_history.add_tool_call(
tool_widget = await chat_history.add_tool_call("search", {"query": "test"}) "tool-1", "search", {"query": "test"}
)
assert isinstance(tool_widget, ToolCallWidget) 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-1")
chat_history.mark_tool_complete(tool_widget) assert tool_widget._completed is True
assert tool_widget._complete is True
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_chat_history_can_add_citations(temp_db_path: Path): async def test_chat_history_can_add_citations(temp_db_path: Path):
"""Test that ChatHistory can display inline citations.""" """Test that ChatHistory can display inline citations."""
from haiku.rag.agents.research.models import Citation 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 from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
mock_client = AsyncMock() app, mock_client = _make_app(temp_db_path)
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): with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
app = ChatApp(temp_db_path, read_only=True)
async with app.run_test(): async with app.run_test():
chat_history = app.query_one(ChatHistory) 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) await chat_history.add_citations(test_citations)
# Verify citation widgets were added
citation_widgets = chat_history.query(CitationWidget) citation_widgets = chat_history.query(CitationWidget)
assert len(list(citation_widgets)) == 2 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 @pytest.mark.asyncio
async def test_chat_history_thinking_indicator(temp_db_path: Path): async def test_chat_history_thinking_indicator(temp_db_path: Path):
"""Test that ChatHistory can show and hide thinking indicator.""" """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 from haiku.rag.chat.widgets.chat_history import ChatHistory, ThinkingWidget
mock_client = AsyncMock() app, mock_client = _make_app(temp_db_path)
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): 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: async with app.run_test() as pilot:
chat_history = app.query_one(ChatHistory) chat_history = app.query_one(ChatHistory)
# Show thinking indicator
await chat_history.show_thinking() await chat_history.show_thinking()
thinking = chat_history.query(ThinkingWidget) thinking = chat_history.query(ThinkingWidget)
assert len(list(thinking)) == 1 assert len(list(thinking)) == 1
# Hide thinking indicator (remove is deferred, so pause)
chat_history.hide_thinking() chat_history.hide_thinking()
await pilot.pause() await pilot.pause()
thinking = chat_history.query(ThinkingWidget) thinking = chat_history.query(ThinkingWidget)
@ -247,67 +185,38 @@ async def test_chat_history_thinking_indicator(temp_db_path: Path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_clear_chat_resets_session(temp_db_path: Path): async def test_clear_chat_resets_state(temp_db_path: Path):
"""Test that clearing chat resets the session state.""" """Test that clearing chat resets state and messages."""
from haiku.rag.chat.app import ChatApp
from haiku.rag.chat.widgets.chat_history import ChatHistory from haiku.rag.chat.widgets.chat_history import ChatHistory
mock_client = AsyncMock() app, mock_client = _make_app(temp_db_path)
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): 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: async with app.run_test() as pilot:
chat_history = app.query_one(ChatHistory) chat_history = app.query_one(ChatHistory)
# Add some messages
await chat_history.add_message("user", "Hello") await chat_history.add_message("user", "Hello")
await chat_history.add_message("assistant", "Hi there") await chat_history.add_message("assistant", "Hi there")
assert len(chat_history.messages) == 2 assert len(chat_history.messages) == 2
# Clear chat via action (available through command palette)
await app.action_clear_chat() await app.action_clear_chat()
await pilot.pause() await pilot.pause()
# Verify messages cleared
assert len(chat_history.messages) == 0 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 @pytest.mark.asyncio
async def test_citation_expand_collapse_with_enter(temp_db_path: Path): async def test_citation_expand_collapse_with_enter(temp_db_path: Path):
"""Test that pressing Enter on a focused citation toggles expand/collapse.""" """Test that pressing Enter on a focused citation toggles expand/collapse."""
from haiku.rag.agents.research.models import Citation 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 from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
mock_client = AsyncMock() app, mock_client = _make_app(temp_db_path)
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): 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: async with app.run_test() as pilot:
chat_history = app.query_one(ChatHistory) chat_history = app.query_one(ChatHistory)
# Add a citation
test_citation = Citation( test_citation = Citation(
index=1, index=1,
document_id="doc1", 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]) await chat_history.add_citations([test_citation])
# Get the citation widget
citation_widget = chat_history.query_one(CitationWidget) citation_widget = chat_history.query_one(CitationWidget)
assert citation_widget.collapsed is True assert citation_widget.collapsed is True
# Focus the citation
citation_widget.focus() citation_widget.focus()
await pilot.pause() await pilot.pause()
# Press Enter to expand
await pilot.press("enter") await pilot.press("enter")
await pilot.pause() await pilot.pause()
assert citation_widget.collapsed is False assert citation_widget.collapsed is False
# Press Enter again to collapse
await pilot.press("enter") await pilot.press("enter")
await pilot.pause() await pilot.pause()
assert citation_widget.collapsed is True assert citation_widget.collapsed is True