11b1bfbc94
A request can carry several of this capability's returns — a model can call search twice in one response — and the carrier was identified by message alone, so each of them received the whole capsule. It is selected by message and part now, so exactly one carries it however many share the request. The chat TUI passed its persisted state into the run, so tool synchronisation mutated it in place while the message history was promoted only on success. A cancelled or failed run therefore kept the evidence the tools had recorded and discarded the messages that justified it, and the next question derived its identity from the shorter history: behind the recorded epoch, refused as non-append-only, the conversation unusable until cleared. The run gets a copy, promoted with the messages or not at all. Five decorators had been left attached to a helper by an earlier extraction, which pytest does not collect, so the resume case they carried was silently untested. The wire test covers both resume shapes again, no prompt and deferred results.
427 lines
15 KiB
Python
427 lines
15 KiB
Python
import asyncio
|
|
import uuid
|
|
from collections.abc import Iterable, Sequence
|
|
from copy import deepcopy
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import textual_image.widget # noqa: F401 - import early for renderer detection
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.messages import (
|
|
BinaryContent,
|
|
FunctionToolCallEvent,
|
|
FunctionToolResultEvent,
|
|
PartDeltaEvent,
|
|
PartEndEvent,
|
|
PartStartEvent,
|
|
TextPart,
|
|
TextPartDelta,
|
|
)
|
|
from pydantic_ai.run import AgentRunResultEvent
|
|
from textual.app import App, SystemCommand
|
|
from textual.binding import Binding
|
|
from textual.widgets import Footer, Header
|
|
from textual.worker import Worker
|
|
|
|
from haiku.rag.capabilities._base import RAGCapabilityBase
|
|
from haiku.rag.capabilities.analysis import AnalysisState
|
|
from haiku.rag.capabilities.compaction import create_capability as create_compaction
|
|
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState
|
|
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
|
|
from haiku.rag.chat.widgets.image_select import ImageAdded
|
|
from haiku.rag.chat.widgets.prompt import (
|
|
FlexibleInput,
|
|
PostableTextArea,
|
|
build_user_prompt,
|
|
)
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import get_config
|
|
from haiku.rag.telemetry import configure as configure_telemetry
|
|
|
|
configure_telemetry(service_name="haiku-rag")
|
|
|
|
if TYPE_CHECKING:
|
|
from textual.app import ComposeResult
|
|
|
|
|
|
RAG_STATE_NAMESPACE = "rag"
|
|
ANALYSIS_STATE_NAMESPACE = "analysis"
|
|
|
|
|
|
@dataclass
|
|
class ChatDeps:
|
|
state: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
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,
|
|
capabilities: Sequence[RAGCapabilityBase[Any]],
|
|
read_only: bool = False,
|
|
model: str | None = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.db_path = db_path
|
|
self._capabilities = capabilities
|
|
self.read_only = read_only
|
|
self._model = model
|
|
self.client: HaikuRAG | None = None
|
|
self.config = get_config()
|
|
self._agent: Agent[ChatDeps, 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] = []
|
|
self._images: list[bytes] = []
|
|
# Stable per-launch id for multi-turn model and telemetry correlation.
|
|
self._conversation_id = str(uuid.uuid4())
|
|
|
|
def compose(self) -> "ComposeResult":
|
|
"""Compose the UI layout."""
|
|
yield Header()
|
|
yield ChatHistory(id="chat-history")
|
|
yield FlexibleInput(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,
|
|
)
|
|
|
|
async def on_mount(self) -> None:
|
|
"""Initialize the app when mounted."""
|
|
client = HaikuRAG(
|
|
db_path=self.db_path,
|
|
config=self.config,
|
|
read_only=self.read_only,
|
|
)
|
|
# Assign only after a successful open: on_unmount must not tear down
|
|
# a client whose __aenter__ failed.
|
|
await client.__aenter__()
|
|
self.client = client
|
|
|
|
self._agent = Agent(
|
|
self._model,
|
|
deps_type=ChatDeps,
|
|
instructions=AGENT_PREAMBLE,
|
|
# A chat is multi-turn by definition, so earlier questions are reduced
|
|
# to the evidence they cited rather than carried whole.
|
|
capabilities=[*self._capabilities, create_compaction()],
|
|
)
|
|
self._state = {}
|
|
for capability in self._capabilities:
|
|
self._state[capability.state_namespace] = (
|
|
capability.state_type().model_dump(mode="json")
|
|
)
|
|
|
|
self.query_one(FlexibleInput).focus()
|
|
|
|
async def on_unmount(self) -> None:
|
|
"""Clean up when unmounting."""
|
|
if self.client:
|
|
await self.client.__aexit__(None, None, None)
|
|
|
|
async def on_flexible_input_submitted(self, event: FlexibleInput.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)
|
|
|
|
user_prompt = build_user_prompt(user_message, self._images)
|
|
self._images = []
|
|
|
|
self._is_processing = True
|
|
self.query_one(FlexibleInput).disabled = True
|
|
self._current_worker = self.run_worker(
|
|
self._run_agent(user_prompt), exclusive=True
|
|
)
|
|
|
|
def on_image_added(self, event: ImageAdded) -> None:
|
|
"""Attach a picked image and insert its token into the prompt."""
|
|
self._images.append(event.data)
|
|
prompt = self.query_one(FlexibleInput)
|
|
prompt.insert_at_cursor(f"[Image #{len(self._images)}]")
|
|
prompt.focus()
|
|
self.notify(f"Attached {event.path.name}")
|
|
|
|
async def _run_agent(self, user_prompt: str | list[str | BinaryContent]) -> None:
|
|
"""Run the agent in a background worker."""
|
|
if not self._agent:
|
|
return
|
|
|
|
chat_history = self.query_one(ChatHistory)
|
|
await chat_history.show_thinking()
|
|
|
|
message = None
|
|
# The run gets a copy: state and message history have to advance together.
|
|
# A cancelled or failed run discards its messages, and state that advanced
|
|
# anyway would leave the next question deriving its identity from a shorter
|
|
# history than the evidence already recorded — refused as non-append-only,
|
|
# with the conversation stuck until it is cleared.
|
|
deps = ChatDeps(state=deepcopy(self._state))
|
|
|
|
try:
|
|
async with self._agent.run_stream_events(
|
|
user_prompt,
|
|
message_history=self._messages,
|
|
conversation_id=self._conversation_id,
|
|
deps=deps,
|
|
) as stream:
|
|
async for event in stream:
|
|
if isinstance(event, PartStartEvent) and isinstance(
|
|
event.part, TextPart
|
|
):
|
|
chat_history.hide_thinking()
|
|
message = await chat_history.add_message("assistant")
|
|
if event.part.content:
|
|
await message.append_delta(event.part.content)
|
|
elif isinstance(event, PartDeltaEvent) and isinstance(
|
|
event.delta, TextPartDelta
|
|
):
|
|
if message:
|
|
await message.append_delta(event.delta.content_delta)
|
|
chat_history.scroll_end(animate=False)
|
|
elif isinstance(event, PartEndEvent) and isinstance(
|
|
event.part, TextPart
|
|
):
|
|
if message:
|
|
await message.finish_stream()
|
|
elif isinstance(event, FunctionToolCallEvent):
|
|
part = event.part
|
|
chat_history.hide_thinking()
|
|
await chat_history.add_tool_call(
|
|
part.tool_call_id, part.tool_name
|
|
)
|
|
chat_history.update_tool_args(
|
|
part.tool_call_id, part.args_as_dict()
|
|
)
|
|
await chat_history.show_thinking("Executing tasks...")
|
|
elif isinstance(event, FunctionToolResultEvent):
|
|
chat_history.mark_tool_complete(event.part.tool_call_id)
|
|
elif isinstance(event, AgentRunResultEvent):
|
|
self._messages = event.result.all_messages()
|
|
self._state = deps.state
|
|
chat_history.hide_thinking()
|
|
await self._show_citations_and_programs(chat_history)
|
|
|
|
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(FlexibleInput)
|
|
chat_input.disabled = False
|
|
chat_input.focus()
|
|
|
|
async def _show_citations_and_programs(self, chat_history: "ChatHistory") -> None:
|
|
"""Show citations and programs from capability states after a response."""
|
|
citations = []
|
|
for namespace in (RAG_STATE_NAMESPACE, ANALYSIS_STATE_NAMESPACE):
|
|
state_data = self._state.get(namespace)
|
|
if not state_data:
|
|
continue
|
|
state_type = RAGState if namespace == RAG_STATE_NAMESPACE else AnalysisState
|
|
state = state_type.model_validate(state_data)
|
|
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)
|
|
|
|
if analysis_data := self._state.get(ANALYSIS_STATE_NAMESPACE):
|
|
analysis_state = AnalysisState.model_validate(analysis_data)
|
|
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()
|
|
self._state = {
|
|
capability.state_namespace: capability.state_type().model_dump(mode="json")
|
|
for capability in self._capabilities
|
|
}
|
|
# 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(FlexibleInput).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, PostableTextArea):
|
|
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_ids = citation.chunk_ids or [citation.chunk_id]
|
|
chunks = []
|
|
for cid in chunk_ids:
|
|
chunk = await self.client.get_chunk_by_id(cid)
|
|
if chunk:
|
|
chunks.append(chunk)
|
|
if not chunks:
|
|
return
|
|
|
|
from haiku.rag.inspector.widgets.visual_modal import VisualGroundingModal
|
|
|
|
await self.push_screen(
|
|
VisualGroundingModal(
|
|
chunk=chunks,
|
|
client=self.client,
|
|
refs=citation.doc_item_refs or None,
|
|
)
|
|
)
|
|
|
|
async def action_show_info(self) -> None:
|
|
"""Show database info modal."""
|
|
if not self.client:
|
|
return
|
|
|
|
from haiku.rag.inspector.widgets.info_modal import InfoModal
|
|
|
|
await self.push_screen(InfoModal(self.client, self.db_path))
|
|
|
|
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
|
|
"""Handle citation selection."""
|
|
chat_history = self.query_one(ChatHistory)
|
|
|
|
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
|
|
|
|
doc_filter = build_multi_document_filter(self._document_filter)
|
|
for namespace, state_type in (
|
|
(RAG_STATE_NAMESPACE, RAGState),
|
|
(ANALYSIS_STATE_NAMESPACE, AnalysisState),
|
|
):
|
|
if namespace in self._state:
|
|
state = state_type.model_validate(self._state[namespace])
|
|
state.document_filter = doc_filter
|
|
self._state[namespace] = state.model_dump(mode="json")
|