480 lines
18 KiB
Python
480 lines
18 KiB
Python
import asyncio
|
|
import json
|
|
import uuid
|
|
from collections.abc import Iterable
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import textual_image.widget # noqa: F401 - import early for renderer detection
|
|
from ag_ui.core import (
|
|
ActivitySnapshotEvent,
|
|
AssistantMessage,
|
|
EventType,
|
|
RunAgentInput,
|
|
StateDeltaEvent,
|
|
TextMessageContentEvent,
|
|
ToolCallArgsEvent,
|
|
ToolCallEndEvent,
|
|
ToolCallStartEvent,
|
|
UserMessage,
|
|
)
|
|
from jsonpatch import JsonPatch
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.ui.ag_ui import AGUIAdapter
|
|
from textual.app import App, SystemCommand
|
|
from textual.binding import Binding
|
|
from textual.widgets import Footer, Header, Input
|
|
from textual.worker import Worker
|
|
|
|
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import get_config
|
|
from haiku.rag.skills.analysis import AnalysisState
|
|
from haiku.rag.skills.rag import RAGState, get_agent_preamble
|
|
from haiku.rag.telemetry import configure as configure_telemetry
|
|
from haiku.skills.agent import (
|
|
SkillToolset,
|
|
run_agui_stream,
|
|
)
|
|
from haiku.skills.models import Skill
|
|
from haiku.skills.prompts import build_system_prompt
|
|
|
|
configure_telemetry()
|
|
|
|
if TYPE_CHECKING:
|
|
from textual.app import ComposeResult
|
|
|
|
|
|
RAG_STATE_NAMESPACE = "rag"
|
|
ANALYSIS_STATE_NAMESPACE = "analysis"
|
|
|
|
|
|
class ChatApp(App):
|
|
"""Textual TUI for conversational RAG."""
|
|
|
|
TITLE = "haiku.rag Chat"
|
|
|
|
CSS = """
|
|
Screen {
|
|
layout: grid;
|
|
grid-size: 1 2;
|
|
grid-rows: 1fr auto;
|
|
background: $surface;
|
|
}
|
|
|
|
#chat-history {
|
|
height: 100%;
|
|
}
|
|
|
|
Header {
|
|
background: $primary;
|
|
}
|
|
|
|
Footer {
|
|
background: $surface-darken-1;
|
|
}
|
|
"""
|
|
|
|
BINDINGS = [
|
|
Binding("escape", "focus_input", "Focus Input", show=False),
|
|
]
|
|
|
|
def __init__(
|
|
self,
|
|
db_path: Path,
|
|
skills: list[Skill],
|
|
read_only: bool = False,
|
|
before: datetime | None = None,
|
|
model: str | None = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.db_path = db_path
|
|
self._skills = skills
|
|
self.read_only = read_only
|
|
self.before = before
|
|
self._model = model
|
|
self.client: HaikuRAG | None = None
|
|
self.config = get_config()
|
|
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._current_worker: Worker[None] | None = None
|
|
self._document_filter: list[str] = []
|
|
# Stable per-launch id so multi-turn chats land in one Logfire
|
|
# conversation. AGUIAdapter reads run_input.thread_id and exports it
|
|
# as the `gen_ai.conversation.id` OTel attribute on every agent run.
|
|
self._conversation_id = str(uuid.uuid4())
|
|
|
|
def compose(self) -> "ComposeResult":
|
|
"""Compose the UI layout."""
|
|
yield Header()
|
|
yield ChatHistory(id="chat-history")
|
|
yield Input(placeholder="Ask a question...", id="chat-input")
|
|
yield Footer()
|
|
|
|
def get_system_commands(self, screen: Any) -> Iterable[SystemCommand]:
|
|
"""Add commands to the command palette."""
|
|
yield from super().get_system_commands(screen)
|
|
yield SystemCommand(
|
|
"Clear chat",
|
|
"Clear the chat history and reset session",
|
|
self.action_clear_chat,
|
|
)
|
|
yield SystemCommand(
|
|
"Filter documents",
|
|
"Select documents to filter searches",
|
|
self.action_show_filter,
|
|
)
|
|
yield SystemCommand(
|
|
"Show visual grounding",
|
|
"Show visual grounding for selected citation",
|
|
self.action_show_visual,
|
|
)
|
|
yield SystemCommand(
|
|
"Database info",
|
|
"Show database information",
|
|
self.action_show_info,
|
|
)
|
|
yield SystemCommand(
|
|
"View state",
|
|
"Show the current session state",
|
|
self.action_view_state,
|
|
)
|
|
|
|
async def on_mount(self) -> None:
|
|
"""Initialize the app when mounted."""
|
|
self.client = HaikuRAG(
|
|
db_path=self.db_path,
|
|
config=self.config,
|
|
read_only=self.read_only,
|
|
before=self.before,
|
|
)
|
|
await self.client.__aenter__()
|
|
|
|
self._toolset = SkillToolset(skills=self._skills)
|
|
self._agent = Agent(
|
|
self._model,
|
|
instructions=build_system_prompt(
|
|
self._toolset.skill_catalog,
|
|
preamble=get_agent_preamble(self.config),
|
|
),
|
|
toolsets=[self._toolset],
|
|
)
|
|
self._state = self._toolset.build_state_snapshot()
|
|
|
|
self.query_one(Input).focus()
|
|
|
|
async def on_unmount(self) -> None:
|
|
"""Clean up when unmounting."""
|
|
if self.client:
|
|
await self.client.__aexit__(None, None, None)
|
|
|
|
async def 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
|
|
|
|
event.input.clear()
|
|
|
|
chat_history = self.query_one(ChatHistory)
|
|
await chat_history.add_message("user", user_message)
|
|
|
|
self._messages.append(
|
|
UserMessage(
|
|
id=str(uuid.uuid4()),
|
|
role="user",
|
|
content=user_message,
|
|
)
|
|
)
|
|
|
|
self._is_processing = True
|
|
self.query_one(Input).disabled = True
|
|
self._current_worker = self.run_worker(
|
|
self._run_agent(user_message), exclusive=True
|
|
)
|
|
|
|
async def _run_agent(self, user_message: str) -> None:
|
|
"""Run the agent in a background worker."""
|
|
if not self._agent or not self._toolset:
|
|
return
|
|
|
|
chat_history = self.query_one(ChatHistory)
|
|
await chat_history.show_thinking()
|
|
|
|
run_input = RunAgentInput(
|
|
thread_id=self._conversation_id,
|
|
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 = ""
|
|
tool_args_deltas: dict[str, str] = {}
|
|
|
|
try:
|
|
async with run_agui_stream(adapter, toolset=self._toolset) as stream:
|
|
async for event in stream:
|
|
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:
|
|
await message.append_delta(event.delta)
|
|
chat_history.scroll_end(animate=False)
|
|
elif event.type == EventType.TEXT_MESSAGE_END:
|
|
if message:
|
|
await message.finish_stream()
|
|
self._messages.append(
|
|
AssistantMessage(
|
|
id=str(uuid.uuid4()),
|
|
role="assistant",
|
|
content=accumulated_text,
|
|
)
|
|
)
|
|
await self._show_citations_and_programs(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
|
|
)
|
|
tool_args_deltas[event.tool_call_id] = ""
|
|
await chat_history.show_thinking("Executing tasks...")
|
|
elif event.type == EventType.TOOL_CALL_ARGS:
|
|
assert isinstance(event, ToolCallArgsEvent)
|
|
tool_args_deltas[event.tool_call_id] = (
|
|
tool_args_deltas.get(event.tool_call_id, "") + event.delta
|
|
)
|
|
try:
|
|
args = json.loads(tool_args_deltas[event.tool_call_id])
|
|
chat_history.update_tool_args(event.tool_call_id, args)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
elif event.type == EventType.TOOL_CALL_END:
|
|
assert isinstance(event, ToolCallEndEvent)
|
|
chat_history.mark_tool_complete(event.tool_call_id)
|
|
elif event.type == EventType.ACTIVITY_SNAPSHOT:
|
|
assert isinstance(event, ActivitySnapshotEvent)
|
|
content = event.content
|
|
if event.activity_type == "skill_tool_call":
|
|
tool_call_id = content["tool_call_id"]
|
|
skill_name = content.get("skill", "")
|
|
tool_name = content["tool_name"]
|
|
display_name = (
|
|
f"{skill_name} → {tool_name}"
|
|
if skill_name
|
|
else tool_name
|
|
)
|
|
args_str = content.get("args", "{}")
|
|
chat_history.hide_thinking()
|
|
await chat_history.add_tool_call(tool_call_id, display_name)
|
|
try:
|
|
args = json.loads(args_str)
|
|
chat_history.update_tool_args(tool_call_id, args)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
await chat_history.show_thinking("Working...")
|
|
elif event.activity_type == "skill_tool_result":
|
|
tool_call_id = content["tool_call_id"]
|
|
chat_history.mark_tool_complete(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()
|
|
if message:
|
|
await message.finish_stream()
|
|
await chat_history.add_message("assistant", "*Cancelled*")
|
|
except Exception as e:
|
|
chat_history.hide_thinking()
|
|
if message:
|
|
await message.finish_stream()
|
|
await chat_history.add_message("assistant", f"Error: {e}")
|
|
finally:
|
|
self._is_processing = False
|
|
self._current_worker = None
|
|
chat_input = self.query_one(Input)
|
|
chat_input.disabled = False
|
|
chat_input.focus()
|
|
|
|
async def _show_citations_and_programs(self, chat_history: "ChatHistory") -> None:
|
|
"""Show citations and programs from skill states after an agent response."""
|
|
if not self._toolset:
|
|
return
|
|
citations = []
|
|
for namespace in (RAG_STATE_NAMESPACE, ANALYSIS_STATE_NAMESPACE):
|
|
state = self._toolset.get_namespace(namespace)
|
|
if not state:
|
|
continue
|
|
cited_ids = getattr(state, "citations", [])
|
|
citation_index = getattr(state, "citation_index", {})
|
|
for cid in cited_ids:
|
|
if cid in citation_index:
|
|
citations.append(citation_index[cid])
|
|
if not citations:
|
|
return
|
|
|
|
picture_bytes: dict[str, list[bytes]] = {}
|
|
if self.client is not None:
|
|
for citation in citations:
|
|
refs = list(citation.picture_refs or [])
|
|
if not refs:
|
|
continue
|
|
blobs: list[bytes] = []
|
|
for ref in refs:
|
|
data = await self.client.document_item_repository.get_picture_bytes(
|
|
citation.document_id, ref
|
|
)
|
|
if data:
|
|
blobs.append(data)
|
|
if blobs:
|
|
picture_bytes[citation.chunk_id] = blobs
|
|
|
|
await chat_history.add_citations(citations, picture_bytes=picture_bytes)
|
|
|
|
analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
|
|
if analysis_state:
|
|
executions = getattr(analysis_state, "executions", [])
|
|
successful = [e for e in executions if e.success]
|
|
if successful:
|
|
await chat_history.add_program(successful[-1].code)
|
|
|
|
async def action_clear_chat(self) -> None:
|
|
"""Clear the chat history and reset session."""
|
|
chat_history = self.query_one(ChatHistory)
|
|
await chat_history.clear_messages()
|
|
self._messages.clear()
|
|
# Reset state
|
|
if self._toolset:
|
|
self._state = self._toolset.build_state_snapshot()
|
|
# Cleared chat starts a fresh Logfire conversation.
|
|
self._conversation_id = str(uuid.uuid4())
|
|
|
|
def action_focus_input(self) -> None:
|
|
"""Focus the input field, or cancel if processing."""
|
|
if self._is_processing and self._current_worker:
|
|
self._current_worker.cancel()
|
|
self.query_one(Input).focus()
|
|
|
|
def _clear_citation_selection(self) -> None:
|
|
"""Clear citation selection."""
|
|
chat_history = self.query_one(ChatHistory)
|
|
for widget in chat_history.query(CitationWidget):
|
|
widget.remove_class("selected")
|
|
|
|
def on_descendant_focus(self, _event: object) -> None:
|
|
"""Clear citation selection when chat input is focused."""
|
|
if isinstance(self.focused, Input) and self.focused.id == "chat-input":
|
|
self._clear_citation_selection()
|
|
|
|
async def action_show_visual(self) -> None:
|
|
"""Show visual grounding for the selected citation."""
|
|
if not self.client:
|
|
return
|
|
|
|
chat_history = self.query_one(ChatHistory)
|
|
selected_widgets = list(chat_history.query(CitationWidget).filter(".selected"))
|
|
if not selected_widgets:
|
|
return
|
|
|
|
citation = selected_widgets[0].citation
|
|
chunk = await self.client.get_chunk_by_id(citation.chunk_id)
|
|
if not chunk:
|
|
return
|
|
|
|
from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal
|
|
|
|
await self.push_screen(VisualGroundingModal(chunk=chunk, client=self.client))
|
|
|
|
async def action_show_info(self) -> None:
|
|
"""Show database info modal."""
|
|
if not self.client:
|
|
return
|
|
|
|
from haiku.rag.inspector.widgets.info_modal import InfoModal
|
|
|
|
await self.push_screen(InfoModal(self.client, self.db_path))
|
|
|
|
def action_view_state(self) -> None:
|
|
"""Show the current session state."""
|
|
from haiku.skills.chat.app import StateScreen
|
|
|
|
self.push_screen(StateScreen(self._state, on_save=self._apply_state_edit))
|
|
|
|
def _apply_state_edit(self, new_state: dict[str, Any]) -> None:
|
|
if self._toolset is None:
|
|
return
|
|
if not isinstance(new_state, dict):
|
|
raise ValueError("state must be a JSON object")
|
|
for namespace, data in new_state.items():
|
|
current = self._toolset.get_namespace(namespace)
|
|
if current is not None:
|
|
type(current).model_validate(data)
|
|
self._toolset.restore_state_snapshot(new_state)
|
|
self._state = self._toolset.build_state_snapshot()
|
|
|
|
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
|
|
"""Handle citation selection."""
|
|
chat_history = self.query_one(ChatHistory)
|
|
|
|
for widget in chat_history.query(CitationWidget):
|
|
widget.remove_class("selected")
|
|
|
|
event.widget.add_class("selected")
|
|
|
|
async def action_show_filter(self) -> None:
|
|
"""Show document filter modal."""
|
|
if not self.client:
|
|
return
|
|
|
|
from haiku.rag.chat.widgets.document_filter_modal import DocumentFilterModal
|
|
|
|
await self.push_screen(
|
|
DocumentFilterModal(
|
|
client=self.client,
|
|
selected=self._document_filter,
|
|
)
|
|
)
|
|
|
|
def on_document_filter_modal_filter_changed(self, event: Any) -> None:
|
|
"""Handle document filter changes from modal."""
|
|
from haiku.rag.tools.filters import build_multi_document_filter
|
|
|
|
self._document_filter = event.selected
|
|
|
|
if self._toolset:
|
|
doc_filter = build_multi_document_filter(self._document_filter)
|
|
rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE)
|
|
if isinstance(rag_state, RAGState):
|
|
rag_state.document_filter = doc_filter
|
|
analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
|
|
if isinstance(analysis_state, AnalysisState):
|
|
analysis_state.document_filter = doc_filter
|
|
self._state = self._toolset.build_state_snapshot()
|