bb92ccf67e
Chat answers with the same capabilities `ask` does, so it federates as naturally as `ask` and `analyze` — but it went through the one-database guard and refused a configured set outright, which left no way to chat across several databases. The guard was the visible half. `run_chat` also defaulted `db_path` to the single default path whenever it was None, so lifting the refusal alone would still have opened one database. It now leaves the path unresolved when `lancedb.databases` names the set, and the client resolves it. Listing and counting documents fan out over the set, which is what the document filter reads, and visual grounding resolves the database holding the cited chunk through the citation's source: chunks, pages and bounding boxes all come from that one database. A limit on a listing means that many documents in total, not that many per database. The info modal reports every database it covers, each under its configured name and without its location, since names are the only identity that leaves the configuration. `database_lines` is what one database reports about itself, shared by both paths, and it reports a failure as a line so one unreachable database does not cost the report on the others. `inspect` stays a one-database command. It browses one database's documents and chunks, so a set has nothing to show it.
431 lines
16 KiB
Python
431 lines
16 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 | None,
|
|
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)
|
|
for cid in state.citations:
|
|
if cid in state.citation_index:
|
|
citations.append(state.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.get_picture_bytes(
|
|
citation.document_id, ref, citation.source
|
|
)
|
|
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)
|
|
successful = [e for e in analysis_state.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
|
|
# Chunks, pages and bounding boxes all come from the database holding the
|
|
# cited chunk. A client covering a set has no repositories of its own.
|
|
client = self.client
|
|
if client._federated:
|
|
if citation.source is None:
|
|
return
|
|
(client,) = await client.clients_for([citation.source])
|
|
chunk_ids = citation.chunk_ids or [citation.chunk_id]
|
|
chunks = []
|
|
for cid in chunk_ids:
|
|
chunk = await 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=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")
|