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

Simplify TUI app

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-11 15:10:50 +02:00
parent 9e63ff1ff7
commit 0434cd068b
No known key found for this signature in database
3 changed files with 55 additions and 441 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -6,8 +6,6 @@ from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any
import jsonpatch
from ag_ui.core import EventType
from pydantic_ai import (
Agent,
AgentStreamEvent,
@ -22,10 +20,6 @@ from haiku.rag.agents.chat.agent import (
create_chat_agent,
trigger_background_summarization,
)
from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatSessionState,
)
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_config
from haiku.rag.tools.context import ToolContext
@ -104,13 +98,11 @@ class ChatApp(App):
self.client: HaikuRAG | None = None
self.config = get_config()
self.agent: Agent[ChatDeps, str] | None = None
self.session_state = ChatSessionState()
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] = []
self._agui_state_snapshot: dict[str, Any] = {}
def compose(self) -> "ComposeResult":
"""Compose the UI layout."""
@ -162,17 +154,11 @@ class ChatApp(App):
self.tool_context = ToolContext()
self.agent = create_chat_agent(self.config, self.client, self.tool_context)
# Initialize session state in tool context
# 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
# Keep ChatSessionState for UI state sync (used by _sync_session_state)
self.session_state = ChatSessionState(
initial_context=self._initial_context,
document_filter=self._document_filter,
)
# Focus the input field
self.query_one(Input).focus()
@ -181,40 +167,6 @@ class ChatApp(App):
if self.client:
await self.client.__aexit__(None, None, None)
def _sync_session_state(self, chat_state: dict[str, Any]) -> None:
"""Sync session_state from AG-UI state.
Updates existing session_state with fields from incoming state,
preserving fields not present in the update (e.g., initial_context).
"""
from haiku.rag.agents.research.models import Citation
# Update specific fields rather than replacing the entire state
if "document_filter" in chat_state:
self.session_state.document_filter = chat_state["document_filter"]
if "citation_registry" in chat_state:
self.session_state.citation_registry = chat_state["citation_registry"]
if "citations" in chat_state:
self.session_state.citations = [
Citation(**c) if isinstance(c, dict) else c
for c in chat_state["citations"]
]
if "qa_history" in chat_state:
from haiku.rag.tools.qa import QAHistoryEntry
self.session_state.qa_history = [
QAHistoryEntry(**qa) if isinstance(qa, dict) else qa
for qa in chat_state["qa_history"]
]
if "session_context" in chat_state:
from haiku.rag.agents.chat.state import SessionContext
ctx = chat_state["session_context"]
if ctx is not None:
self.session_state.session_context = (
SessionContext(**ctx) if isinstance(ctx, dict) else ctx
)
async def _handle_stream_event(self, event: AgentStreamEvent) -> None:
"""Handle streaming events from the agent."""
chat_history = self.query_one(ChatHistory)
@ -232,32 +184,6 @@ class ChatApp(App):
widget = self._tool_call_widgets[tool_call_id]
chat_history.mark_tool_complete(widget)
# Extract citations from state events 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 not hasattr(meta_event, "type"):
continue
if meta_event.type == EventType.STATE_SNAPSHOT:
snapshot = getattr(meta_event, "snapshot", {})
self._agui_state_snapshot = snapshot
chat_state = snapshot.get(AGUI_STATE_KEY, snapshot)
self._sync_session_state(chat_state)
elif meta_event.type == EventType.STATE_DELTA:
delta = getattr(meta_event, "delta", [])
if delta:
patch = jsonpatch.JsonPatch(delta)
self._agui_state_snapshot = patch.apply(
self._agui_state_snapshot
)
chat_state = self._agui_state_snapshot.get(
AGUI_STATE_KEY, self._agui_state_snapshot
)
self._sync_session_state(chat_state)
async def _event_stream_handler(
self,
_ctx: RunContext[ChatDeps],
@ -290,7 +216,9 @@ class ChatApp(App):
# Clear for new query
self._tool_call_widgets.clear()
self.session_state.citations.clear()
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
if session_state:
session_state.citations.clear()
# Run agent in a worker to keep UI responsive
self._is_processing = True
@ -310,39 +238,20 @@ class ChatApp(App):
await chat_history.show_thinking()
try:
# Initialize AGUI state snapshot from session state for delta application
if self.session_state:
self._agui_state_snapshot = {
AGUI_STATE_KEY: self.session_state.model_dump(mode="json")
}
# Sync session state to tool context before running
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
if session_state is not None:
session_state.document_filter = self.session_state.document_filter
session_state.citation_registry = self.session_state.citation_registry
session_state.citations = list(self.session_state.citations)
# Sync initial_context to QA session state
# 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 (
not qa_session_state.session_context
and self.session_state.initial_context
):
qa_session_state.session_context = (
self.session_state.initial_context
)
if not qa_session_state.session_context and self._initial_context:
qa_session_state.session_context = self._initial_context
deps = ChatDeps(
config=self.config,
tool_context=self.tool_context,
is_new=False,
state_key=AGUI_STATE_KEY,
)
async with self.agent.run_stream(
@ -367,27 +276,14 @@ class ChatApp(App):
# Update message history with this conversation
self._message_history = stream.all_messages()
# Add citations captured from tool metadata
if self.session_state.citations:
await chat_history.add_citations(self.session_state.citations)
# 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)
# Sync session context from QASessionState to ChatSessionState for modal
qa_session_state = self.tool_context.get(
QA_SESSION_NAMESPACE, QASessionState
)
if qa_session_state is not None and qa_session_state.session_context:
from datetime import datetime
from haiku.rag.agents.chat.state import SessionContext
self.session_state.session_context = SessionContext(
summary=qa_session_state.session_context,
last_updated=datetime.now(),
)
except asyncio.CancelledError:
chat_history.hide_thinking()
await chat_history.add_message("assistant", "*Cancelled*")
@ -403,16 +299,18 @@ class ChatApp(App):
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._agui_state_snapshot = {}
# Reset context lock and session state (reset to CLI value)
self._context_locked = False
self.session_state = ChatSessionState(
initial_context=self._initial_context,
document_filter=self._document_filter,
# 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())
def action_focus_input(self) -> None:
"""Focus the input field, or cancel if processing."""
@ -462,29 +360,30 @@ class ChatApp(App):
async def action_show_context(self) -> None:
"""Show context modal (edit initial context or view session context)."""
from haiku.rag.agents.chat.state import SessionContext
from haiku.rag.chat.widgets.context_modal import ContextModal
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
# Sync session context from QASessionState before showing modal
session_context = None
qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState)
if qa_session_state is not None and qa_session_state.session_context:
from datetime import datetime
from haiku.rag.agents.chat.state import SessionContext
self.session_state.session_context = SessionContext(
if qa_session_state and qa_session_state.session_context:
session_context = SessionContext(
summary=qa_session_state.session_context,
last_updated=datetime.now(),
)
await self.push_screen(
ContextModal(self.session_state, is_locked=self._context_locked)
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 self.session_state and not self._context_locked:
self.session_state.initial_context = event.context or None
if not self._context_locked:
self._initial_context = event.context or None
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
"""Handle citation selection."""
@ -514,5 +413,6 @@ 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
if self.session_state:
self.session_state.document_filter = self._document_filter
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
if session_state:
session_state.document_filter = self._document_filter

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

@ -8,7 +8,7 @@ from textual.screen import ModalScreen
from textual.widgets import Button, Markdown, Static, TextArea
if TYPE_CHECKING:
from haiku.rag.agents.chat.state import ChatSessionState
from haiku.rag.agents.chat.state import SessionContext
class ContextModal(ModalScreen): # pragma: no cover
@ -80,20 +80,20 @@ class ContextModal(ModalScreen): # pragma: no cover
self.context = context
def __init__(
self, session_state: "ChatSessionState | None", is_locked: bool = False
self,
initial_context: str | None = None,
session_context: "SessionContext | None" = None,
is_locked: bool = False,
) -> None:
super().__init__()
self.session_state = session_state
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_state
and self.session_state.session_context
and self.session_state.session_context.summary
)
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:
@ -105,9 +105,7 @@ class ContextModal(ModalScreen): # pragma: no cover
"This will be locked after you send your first message.",
id="context-description",
)
initial_value = ""
if self.session_state and self.session_state.initial_context:
initial_value = self.session_state.initial_context
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")
@ -124,13 +122,10 @@ class ContextModal(ModalScreen): # pragma: no cover
yield Button("Close", id="cancel-btn", variant="primary")
def _get_session_content(self) -> str:
if not self.session_state:
return "*No session state.*"
if not self.session_state.session_context:
if not self._session_context:
return "*No session context yet. Ask a question first.*"
ctx = self.session_state.session_context
ctx = self._session_context
updated = (
ctx.last_updated.strftime("%Y-%m-%d %H:%M:%S")
if ctx.last_updated

305
tests/chat/test_chat_app.py
View file

@ -5,6 +5,7 @@ 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()
@ -273,302 +274,20 @@ async def test_clear_chat_resets_session(temp_db_path: Path):
# Verify messages cleared
assert len(chat_history.messages) == 0
# Verify session state reset
assert app.session_state is not None
assert app.session_state.qa_history == []
assert app.session_state.citations == []
# 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 == {}
@pytest.mark.asyncio
async def test_handle_stream_event_extracts_citations_from_state_snapshot(
temp_db_path: Path,
):
"""Test that _handle_stream_event extracts citations from STATE_SNAPSHOT events."""
from ag_ui.core import EventType, StateSnapshotEvent
from pydantic_ai import FunctionToolResultEvent
from pydantic_ai.messages import ToolReturnPart
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
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():
# Create a STATE_SNAPSHOT event with citations
snapshot_event = StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot={
AGUI_STATE_KEY: {
"citations": [
{
"index": 1,
"document_id": "doc1",
"chunk_id": "chunk1",
"document_uri": "test.pdf",
"document_title": "Test Doc",
"content": "Test content",
}
],
"qa_history": [],
"citation_registry": {"chunk1": 1},
}
},
)
# Create a tool result with the snapshot in metadata
tool_return = ToolReturnPart(
tool_name="search",
content="Found results",
tool_call_id="test-call-1",
metadata=[snapshot_event],
)
event = FunctionToolResultEvent(result=tool_return)
# Handle the event
await app._handle_stream_event(event)
# Session state should be synced
assert len(app.session_state.citations) == 1
assert app.session_state.citations[0].chunk_id == "chunk1"
assert app.session_state.citation_registry == {"chunk1": 1}
@pytest.mark.asyncio
async def test_handle_stream_event_extracts_citations_from_state_delta(
temp_db_path: Path,
):
"""Test that _handle_stream_event extracts citations from STATE_DELTA events.
This test demonstrates that state deltas need to be applied to extract citations.
After the first STATE_SNAPSHOT, subsequent tool calls emit STATE_DELTA events
containing JSON Patch operations.
"""
from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
from pydantic_ai import FunctionToolResultEvent
from pydantic_ai.messages import ToolReturnPart
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
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():
# First, handle a STATE_SNAPSHOT to establish initial state
initial_snapshot = StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot={
AGUI_STATE_KEY: {
"citations": [],
"qa_history": [],
"citation_registry": {},
}
},
)
tool_return1 = ToolReturnPart(
tool_name="search",
content="Initial search",
tool_call_id="test-call-1",
metadata=[initial_snapshot],
)
event1 = FunctionToolResultEvent(result=tool_return1)
await app._handle_stream_event(event1)
assert len(app.session_state.citations) == 0
# Now handle a STATE_DELTA event that adds citations
delta_event = StateDeltaEvent(
type=EventType.STATE_DELTA,
delta=[
{
"op": "replace",
"path": f"/{AGUI_STATE_KEY}/citations",
"value": [
{
"index": 1,
"document_id": "doc1",
"chunk_id": "chunk1",
"document_uri": "test.pdf",
"document_title": "Test Doc",
"content": "Test content from delta",
}
],
},
{
"op": "add",
"path": f"/{AGUI_STATE_KEY}/citation_registry/chunk1",
"value": 1,
},
],
)
tool_return2 = ToolReturnPart(
tool_name="ask",
content="Answer with citations",
tool_call_id="test-call-2",
metadata=[delta_event],
)
event2 = FunctionToolResultEvent(result=tool_return2)
# Handle the delta event
await app._handle_stream_event(event2)
# Session state should be synced
assert len(app.session_state.citations) == 1
assert app.session_state.citations[0].chunk_id == "chunk1"
assert app.session_state.citations[0].content == "Test content from delta"
@pytest.mark.asyncio
async def test_handle_stream_event_delta_with_preinitialized_state(
temp_db_path: Path,
):
"""Test delta handling when _agui_state_snapshot is pre-initialized.
This is the actual TUI scenario: session_state exists from the start,
so the agent emits deltas (not snapshots) even on the first tool call.
The TUI pre-initializes _agui_state_snapshot from session_state.
"""
from ag_ui.core import EventType, StateDeltaEvent
from pydantic_ai import FunctionToolResultEvent
from pydantic_ai.messages import ToolReturnPart
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
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():
# Pre-initialize _agui_state_snapshot (simulating what _run_agent does)
app._agui_state_snapshot = {
AGUI_STATE_KEY: {
"citations": [],
"qa_history": [],
"citation_registry": {},
"document_filter": [],
"initial_context": None,
"session_context": None,
}
}
# Now handle a STATE_DELTA event directly (no prior snapshot event)
delta_event = StateDeltaEvent(
type=EventType.STATE_DELTA,
delta=[
{
"op": "replace",
"path": f"/{AGUI_STATE_KEY}/citations",
"value": [
{
"index": 1,
"document_id": "doc1",
"chunk_id": "chunk1",
"document_uri": "test.pdf",
"document_title": "Test Doc",
"content": "Content from first delta",
}
],
},
{
"op": "add",
"path": f"/{AGUI_STATE_KEY}/citation_registry/chunk1",
"value": 1,
},
],
)
tool_return = ToolReturnPart(
tool_name="ask",
content="Answer with citations",
tool_call_id="test-call-1",
metadata=[delta_event],
)
event = FunctionToolResultEvent(result=tool_return)
# Handle the delta event
await app._handle_stream_event(event)
# Session state should be synced
assert len(app.session_state.citations) == 1
assert app.session_state.citations[0].chunk_id == "chunk1"
assert app.session_state.citations[0].content == "Content from first delta"
@pytest.mark.asyncio
async def test_handle_stream_event_syncs_session_context(temp_db_path: Path):
"""Test that _handle_stream_event syncs session_context to session_state."""
from ag_ui.core import EventType, StateDeltaEvent
from pydantic_ai import FunctionToolResultEvent
from pydantic_ai.messages import ToolReturnPart
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
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():
# Pre-initialize state
app._agui_state_snapshot = {
AGUI_STATE_KEY: {
"citations": [],
"qa_history": [],
"citation_registry": {},
"document_filter": [],
"initial_context": None,
"session_context": None,
}
}
# Verify session_context starts as None
assert app.session_state.session_context is None
# Handle a delta that adds session_context
delta_event = StateDeltaEvent(
type=EventType.STATE_DELTA,
delta=[
{
"op": "replace",
"path": f"/{AGUI_STATE_KEY}/session_context",
"value": {
"summary": "User asked about Python async patterns.",
"last_updated": "2025-01-15T10:30:00",
},
},
],
)
tool_return = ToolReturnPart(
tool_name="ask",
content="Answer about async",
tool_call_id="test-call-1",
metadata=[delta_event],
)
event = FunctionToolResultEvent(result=tool_return)
await app._handle_stream_event(event)
# Session context should be synced to session_state
assert app.session_state.session_context is not None
assert (
app.session_state.session_context.summary
== "User asked about Python async patterns."
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