Rewrite chat TUI and app backend with haiku.skills
This commit is contained in:
Yiorgis Gozadinos
parent
8176103eb0
commit
643ed20d6f
7 changed files with 289 additions and 477 deletions
|
|
@ -3,8 +3,9 @@ import os
|
|||
from pathlib import Path
|
||||
|
||||
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.ag_ui import AGUIAdapter
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
|
|
@ -12,22 +13,14 @@ from starlette.requests import Request
|
|||
from starlette.responses import JSONResponse, Response, StreamingResponse
|
||||
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.config import load_yaml_config
|
||||
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))
|
||||
|
||||
# Cache ToolContext instances by thread_id across requests
|
||||
context_cache = ToolContextCache()
|
||||
|
||||
# Configure logfire (only sends data if LOGFIRE_TOKEN is present)
|
||||
try:
|
||||
import logfire
|
||||
|
|
@ -69,34 +62,24 @@ def get_client() -> HaikuRAG:
|
|||
return _client
|
||||
|
||||
|
||||
# Toolkit and agent are created once at module level
|
||||
chat_toolkit = build_chat_toolkit(Config)
|
||||
agent = create_chat_agent(Config, toolkit=chat_toolkit)
|
||||
# Create skill, toolset, and agent
|
||||
skill = create_skill(db_path=db_path, config=Config)
|
||||
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:
|
||||
"""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.
|
||||
"""
|
||||
"""Chat streaming endpoint with AG-UI protocol."""
|
||||
body = await request.body()
|
||||
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
|
||||
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)
|
||||
event_stream = adapter.run_stream(deps=deps)
|
||||
event_stream = adapter.run_stream()
|
||||
sse_event_stream = adapter.encode_stream(event_stream)
|
||||
|
||||
return StreamingResponse(
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ def run_chat(
|
|||
db_path: Path | None = None,
|
||||
read_only: bool = False,
|
||||
before: datetime | None = None,
|
||||
initial_context: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> None:
|
||||
"""Run the chat TUI.
|
||||
|
||||
|
|
@ -14,7 +14,7 @@ def run_chat(
|
|||
db_path: Path to the LanceDB database. If None, uses default from config.
|
||||
read_only: Whether to open the database in read-only mode.
|
||||
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:
|
||||
from haiku.rag.chat.app import ChatApp
|
||||
|
|
@ -24,15 +24,19 @@ def run_chat(
|
|||
) from e
|
||||
|
||||
from haiku.rag.config import get_config
|
||||
from haiku.rag.skills.rag import create_skill
|
||||
|
||||
config = get_config()
|
||||
if db_path is None:
|
||||
db_path = config.storage.data_dir / "haiku.rag.lancedb"
|
||||
|
||||
skill = create_skill(db_path=db_path, config=config)
|
||||
|
||||
app = ChatApp(
|
||||
db_path,
|
||||
skill=skill,
|
||||
read_only=read_only,
|
||||
before=before,
|
||||
initial_context=initial_context,
|
||||
model=model,
|
||||
)
|
||||
app.run()
|
||||
|
|
|
|||
|
|
@ -1,30 +1,15 @@
|
|||
# pyright: reportPossiblyUnboundVariable=false
|
||||
import asyncio
|
||||
import uuid
|
||||
from collections.abc import AsyncIterable, Iterable
|
||||
from collections.abc import Iterable
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
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.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:
|
||||
from textual.app import ComposeResult
|
||||
|
|
@ -39,6 +24,20 @@ except ImportError: # pragma: no cover
|
|||
|
||||
try:
|
||||
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.binding import Binding
|
||||
from textual.widgets import Footer, Header, Input
|
||||
|
|
@ -53,6 +52,9 @@ except ImportError: # pragma: no cover
|
|||
SystemCommand = object # type: ignore
|
||||
|
||||
|
||||
RAG_STATE_NAMESPACE = "rag"
|
||||
|
||||
|
||||
class ChatApp(App):
|
||||
"""Textual TUI for conversational RAG."""
|
||||
|
||||
|
|
@ -86,23 +88,25 @@ class ChatApp(App):
|
|||
def __init__(
|
||||
self,
|
||||
db_path: Path,
|
||||
skill: Skill,
|
||||
read_only: bool = False,
|
||||
before: datetime | None = None,
|
||||
initial_context: str | None = None,
|
||||
model: str | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.db_path = db_path
|
||||
self._skill = skill
|
||||
self.read_only = read_only
|
||||
self.before = before
|
||||
self._initial_context = initial_context
|
||||
self._context_locked = False
|
||||
self._model = model or "openai:gpt-4o"
|
||||
self.client: HaikuRAG | None = None
|
||||
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._tool_call_widgets: dict[str, Any] = {}
|
||||
self._current_worker: Worker[None] | None = None
|
||||
self._message_history: list[ModelMessage] = []
|
||||
self._document_filter: list[str] = []
|
||||
|
||||
def compose(self) -> "ComposeResult":
|
||||
|
|
@ -136,9 +140,9 @@ class ChatApp(App):
|
|||
self.action_show_info,
|
||||
)
|
||||
yield SystemCommand(
|
||||
"Memory",
|
||||
"View/edit context (editable before first message)",
|
||||
self.action_show_context,
|
||||
"View state",
|
||||
"Show the current session state",
|
||||
self.action_view_state,
|
||||
)
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
|
|
@ -151,17 +155,14 @@ class ChatApp(App):
|
|||
)
|
||||
await self.client.__aenter__()
|
||||
|
||||
# Create toolkit, context, and agent
|
||||
self.toolkit = build_chat_toolkit(self.config)
|
||||
self.tool_context = self.toolkit.create_context(state_key=AGUI_STATE_KEY)
|
||||
self.agent = create_chat_agent(self.config, toolkit=self.toolkit)
|
||||
self._toolset = SkillToolset(skills=[self._skill])
|
||||
self._agent = Agent(
|
||||
self._model,
|
||||
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()
|
||||
|
||||
async def on_unmount(self) -> None:
|
||||
|
|
@ -169,60 +170,25 @@ class ChatApp(App):
|
|||
if self.client:
|
||||
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:
|
||||
"""Handle user input submission."""
|
||||
user_message = event.value.strip()
|
||||
if not user_message or self._is_processing:
|
||||
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()
|
||||
|
||||
# Add user message to history
|
||||
chat_history = self.query_one(ChatHistory)
|
||||
await chat_history.add_message("user", user_message)
|
||||
|
||||
# Clear for new query
|
||||
self._tool_call_widgets.clear()
|
||||
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
|
||||
if session_state:
|
||||
session_state.citations.clear()
|
||||
self._messages.append(
|
||||
UserMessage(
|
||||
id=str(uuid.uuid4()),
|
||||
role="user",
|
||||
content=user_message,
|
||||
)
|
||||
)
|
||||
|
||||
# Run agent in a worker to keep UI responsive
|
||||
self._is_processing = True
|
||||
self.query_one(Input).disabled = True
|
||||
self._current_worker = self.run_worker(
|
||||
|
|
@ -231,64 +197,75 @@ class ChatApp(App):
|
|||
|
||||
async def _run_agent(self, user_message: str) -> None:
|
||||
"""Run the agent in a background worker."""
|
||||
if not self.client or not self.agent:
|
||||
if not self._agent or not self._toolset:
|
||||
return
|
||||
|
||||
chat_history = self.query_one(ChatHistory)
|
||||
|
||||
# Show thinking indicator
|
||||
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:
|
||||
# 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 qa_session_state.session_context is None and self._initial_context:
|
||||
from haiku.rag.tools.session import SessionContext
|
||||
|
||||
qa_session_state.session_context = SessionContext(
|
||||
summary=self._initial_context
|
||||
async for event in adapter.run_stream():
|
||||
if not isinstance(event, BaseEvent):
|
||||
continue
|
||||
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:
|
||||
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,
|
||||
)
|
||||
)
|
||||
|
||||
deps = ChatDeps(
|
||||
config=self.config,
|
||||
client=self.client,
|
||||
tool_context=self.tool_context,
|
||||
)
|
||||
|
||||
async with self.agent.run_stream(
|
||||
user_message,
|
||||
deps=deps,
|
||||
message_history=self._message_history,
|
||||
event_stream_handler=self._event_stream_handler,
|
||||
) as stream:
|
||||
# Hide thinking when we start getting content
|
||||
chat_history.hide_thinking()
|
||||
|
||||
# Create assistant message for streaming
|
||||
assistant_msg = await chat_history.add_message("assistant", "")
|
||||
|
||||
# Stream text updates
|
||||
async for text in stream.stream_text():
|
||||
assistant_msg.update_content(text)
|
||||
chat_history.scroll_end(animate=False)
|
||||
# Yield to event loop to keep UI responsive
|
||||
await asyncio.sleep(0)
|
||||
|
||||
# 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)
|
||||
# Show citations from RAG state
|
||||
await self._show_citations(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
|
||||
)
|
||||
await chat_history.show_thinking("Executing tasks...")
|
||||
elif event.type == EventType.TOOL_CALL_END:
|
||||
assert isinstance(event, ToolCallEndEvent)
|
||||
chat_history.mark_tool_complete(event.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()
|
||||
|
|
@ -303,20 +280,26 @@ class ChatApp(App):
|
|||
chat_input.disabled = False
|
||||
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:
|
||||
"""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._context_locked = False
|
||||
# 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())
|
||||
self._messages.clear()
|
||||
# Reset state
|
||||
if self._toolset:
|
||||
self._state = self._toolset.build_state_snapshot()
|
||||
|
||||
def action_focus_input(self) -> None:
|
||||
"""Focus the input field, or cancel if processing."""
|
||||
|
|
@ -340,7 +323,6 @@ class ChatApp(App):
|
|||
if not self.client:
|
||||
return
|
||||
|
||||
# Get citation from selected widget directly
|
||||
chat_history = self.query_one(ChatHistory)
|
||||
selected_widgets = list(chat_history.query(CitationWidget).filter(".selected"))
|
||||
if not selected_widgets:
|
||||
|
|
@ -364,38 +346,19 @@ class ChatApp(App):
|
|||
|
||||
await self.push_screen(InfoModal(self.client, self.db_path))
|
||||
|
||||
async def action_show_context(self) -> None:
|
||||
"""Show context modal (edit initial context or view session context)."""
|
||||
from haiku.rag.chat.widgets.context_modal import ContextModal
|
||||
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
|
||||
def action_view_state(self) -> None:
|
||||
"""Show the current session state."""
|
||||
from haiku.skills.chat.app import StateScreen
|
||||
|
||||
session_context = None
|
||||
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
|
||||
self.push_screen(StateScreen(self._state))
|
||||
|
||||
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
|
||||
"""Handle citation selection."""
|
||||
chat_history = self.query_one(ChatHistory)
|
||||
|
||||
# Remove selected class from all citations
|
||||
for widget in chat_history.query(CitationWidget):
|
||||
widget.remove_class("selected")
|
||||
|
||||
# Add selected class to the widget that was focused
|
||||
event.widget.add_class("selected")
|
||||
|
||||
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:
|
||||
"""Handle document filter changes from modal."""
|
||||
self._document_filter = event.selected
|
||||
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
|
||||
if session_state:
|
||||
session_state.document_filter = self._document_filter
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from textual.containers import Horizontal, VerticalScroll
|
||||
from textual.message import Message
|
||||
|
|
@ -32,52 +32,48 @@ class ChatMessage(Static):
|
|||
|
||||
|
||||
class ToolCallWidget(Static):
|
||||
"""Styled inline display of a tool call."""
|
||||
"""Displays a single tool call with status indicator."""
|
||||
|
||||
TOOL_LABELS = {
|
||||
"search": "Searching",
|
||||
"ask": "Asking",
|
||||
"get_document": "Fetching",
|
||||
}
|
||||
|
||||
def __init__(self, tool_name: str, args: dict | None = None, **kwargs) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
tool_call_id: str,
|
||||
tool_name: str,
|
||||
args: dict[str, Any] | None = None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.tool_call_id = tool_call_id
|
||||
self.tool_name = tool_name
|
||||
self.args = args or {}
|
||||
self._complete = False
|
||||
self._completed = False
|
||||
|
||||
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"):
|
||||
if self._complete:
|
||||
if self._completed:
|
||||
yield Static("✓", classes="tool-status")
|
||||
else:
|
||||
yield LoadingIndicator(classes="tool-spinner")
|
||||
yield Static(label, classes="tool-badge")
|
||||
yield Static(desc, classes="tool-desc")
|
||||
yield Static(self.tool_name, classes="tool-badge")
|
||||
desc = self._build_description()
|
||||
if desc:
|
||||
yield Static(desc, classes="tool-desc")
|
||||
|
||||
def mark_complete(self) -> None:
|
||||
self._complete = True
|
||||
def _build_description(self) -> str:
|
||||
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)
|
||||
|
||||
|
||||
|
|
@ -102,7 +98,6 @@ class CitationWidget(Collapsible):
|
|||
pages += "..."
|
||||
title += f" (p.{pages})"
|
||||
|
||||
# Build content widgets
|
||||
content = citation.content
|
||||
if len(content) > 500:
|
||||
content = content[:500] + "..."
|
||||
|
|
@ -132,10 +127,22 @@ class CitationWidget(Collapsible):
|
|||
class ThinkingWidget(Static):
|
||||
"""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":
|
||||
with Horizontal(classes="thinking-row"):
|
||||
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):
|
||||
|
|
@ -303,6 +310,7 @@ class ChatHistory(VerticalScroll):
|
|||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.messages: list[tuple[str, str]] = []
|
||||
self._tool_widgets: dict[str, ToolCallWidget] = {}
|
||||
|
||||
async def add_message(self, role: str, content: str = "") -> ChatMessage:
|
||||
"""Add a message to the chat history."""
|
||||
|
|
@ -313,18 +321,24 @@ class ChatHistory(VerticalScroll):
|
|||
return message_widget
|
||||
|
||||
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:
|
||||
"""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)
|
||||
self.scroll_end(animate=False)
|
||||
return widget
|
||||
|
||||
def mark_tool_complete(self, widget: ToolCallWidget) -> None:
|
||||
"""Mark a tool call as complete."""
|
||||
widget.mark_complete()
|
||||
widget.add_class("complete")
|
||||
def mark_tool_complete(self, tool_call_id: str) -> None:
|
||||
"""Mark a tool call as complete by its ID."""
|
||||
widget = self._tool_widgets.get(tool_call_id)
|
||||
if widget:
|
||||
widget.mark_completed()
|
||||
widget.add_class("complete")
|
||||
|
||||
async def add_citations(self, citations: list[Citation]) -> None:
|
||||
"""Add citations inline after a response."""
|
||||
|
|
@ -336,9 +350,12 @@ class ChatHistory(VerticalScroll):
|
|||
await self.mount(widget)
|
||||
self.scroll_end(animate=False)
|
||||
|
||||
async def show_thinking(self) -> None:
|
||||
async def show_thinking(self, text: str = "Thinking...") -> None:
|
||||
"""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)
|
||||
|
||||
def hide_thinking(self) -> None:
|
||||
|
|
@ -351,4 +368,5 @@ class ChatHistory(VerticalScroll):
|
|||
async def clear_messages(self) -> None:
|
||||
"""Clear all messages from the chat history."""
|
||||
self.messages.clear()
|
||||
self._tool_widgets.clear()
|
||||
await self.remove_children()
|
||||
|
|
|
|||
|
|
@ -1,22 +1,12 @@
|
|||
from typing import TYPE_CHECKING
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.containers import Horizontal, Vertical, VerticalScroll
|
||||
from textual.message import Message
|
||||
from textual.screen import ModalScreen
|
||||
from textual.widgets import Button, Markdown, Static, TextArea
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.tools.session import SessionContext
|
||||
from textual.widgets import Button, Markdown, Static
|
||||
|
||||
|
||||
class ContextModal(ModalScreen): # pragma: no cover
|
||||
"""Modal screen for viewing/editing context.
|
||||
|
||||
Before first message (not locked, no session context): Edit initial context
|
||||
After first message (locked or has session context): View session context
|
||||
"""
|
||||
"""Modal screen for viewing session Q&A history."""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "cancel", "Close", show=False),
|
||||
|
|
@ -49,12 +39,6 @@ class ContextModal(ModalScreen): # pragma: no cover
|
|||
color: $text-muted;
|
||||
}
|
||||
|
||||
#context-editor {
|
||||
height: 12;
|
||||
min-height: 8;
|
||||
max-height: 16;
|
||||
}
|
||||
|
||||
#context-content {
|
||||
height: 1fr;
|
||||
max-height: 16;
|
||||
|
|
@ -73,81 +57,39 @@ class ContextModal(ModalScreen): # pragma: no cover
|
|||
}
|
||||
"""
|
||||
|
||||
class ContextUpdated(Message):
|
||||
"""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:
|
||||
def __init__(self, qa_history: list | None = None) -> None:
|
||||
super().__init__()
|
||||
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_context and self._session_context.summary
|
||||
return not self._is_locked and not has_session_context
|
||||
self._qa_history = qa_history or []
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
with Vertical(id="context-container"):
|
||||
if self._is_edit_mode:
|
||||
yield Static("[bold]Initial Context[/bold]", id="context-header")
|
||||
yield Static(
|
||||
"Set background context to guide the conversation. "
|
||||
"This will be locked after you send your first message.",
|
||||
id="context-description",
|
||||
)
|
||||
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")
|
||||
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")
|
||||
yield Static("[bold]Session Context[/bold]", id="context-header")
|
||||
yield Static(
|
||||
"Questions and answers from this session.",
|
||||
id="context-description",
|
||||
)
|
||||
with VerticalScroll(id="context-content"):
|
||||
yield Markdown(self._get_content())
|
||||
with Horizontal(id="button-row"):
|
||||
yield Button("Close", id="cancel-btn", variant="primary")
|
||||
|
||||
def _get_session_content(self) -> str:
|
||||
if not self._session_context:
|
||||
return "*No session context yet. Ask a question first.*"
|
||||
def _get_content(self) -> str:
|
||||
if not self._qa_history:
|
||||
return "*No questions asked yet.*"
|
||||
|
||||
ctx = self._session_context
|
||||
updated = (
|
||||
ctx.last_updated.strftime("%Y-%m-%d %H:%M:%S")
|
||||
if ctx.last_updated
|
||||
else "unknown"
|
||||
)
|
||||
parts = []
|
||||
for entry in self._qa_history:
|
||||
q = getattr(entry, "question", str(entry))
|
||||
a = getattr(entry, "answer", "")
|
||||
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:
|
||||
"""Handle button presses."""
|
||||
if event.button.id == "cancel-btn":
|
||||
self.action_cancel()
|
||||
elif event.button.id == "save-btn":
|
||||
self.action_save()
|
||||
|
||||
def action_cancel(self) -> None:
|
||||
"""Cancel and close without saving."""
|
||||
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))
|
||||
"""Cancel and close."""
|
||||
self.app.pop_screen()
|
||||
|
|
|
|||
|
|
@ -620,10 +620,10 @@ def chat( # pragma: no cover
|
|||
"--db",
|
||||
help="Path to the LanceDB database file",
|
||||
),
|
||||
initial_context: str | None = typer.Option(
|
||||
model: str | None = typer.Option(
|
||||
None,
|
||||
"--initial-context",
|
||||
help="Initial background context to provide to the conversation",
|
||||
"--model",
|
||||
help="Model to use for the chat agent (e.g. openai:gpt-4o)",
|
||||
),
|
||||
):
|
||||
"""Launch the chat TUI for conversational RAG."""
|
||||
|
|
@ -635,7 +635,7 @@ def chat( # pragma: no cover
|
|||
db_path,
|
||||
read_only=_read_only,
|
||||
before=_before,
|
||||
initial_context=initial_context,
|
||||
model=model,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ 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()
|
||||
|
||||
|
|
@ -21,19 +20,46 @@ def test_chat_command():
|
|||
mock_chat.assert_called_once()
|
||||
|
||||
|
||||
@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.app import ChatApp
|
||||
from haiku.rag.chat.widgets.chat_history import ChatHistory
|
||||
|
||||
def _make_mock_client():
|
||||
"""Create a mock HaikuRAG client."""
|
||||
mock_client = AsyncMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
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):
|
||||
app = ChatApp(temp_db_path, read_only=True)
|
||||
|
||||
async with app.run_test():
|
||||
chat_history = app.query_one(ChatHistory)
|
||||
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
|
||||
async def test_chat_app_quit_binding(temp_db_path: Path):
|
||||
"""Test that pressing ctrl+q quits the app."""
|
||||
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)
|
||||
app, mock_client = _make_app(temp_db_path)
|
||||
|
||||
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:
|
||||
# App should be running
|
||||
assert app.is_running
|
||||
|
||||
# Press ctrl+q to quit
|
||||
await pilot.press("ctrl+q")
|
||||
|
||||
# App should have exited
|
||||
assert not app.is_running
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_history_can_add_message(temp_db_path: Path):
|
||||
"""Test that ChatHistory can display messages."""
|
||||
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)
|
||||
app, mock_client = _make_app(temp_db_path)
|
||||
|
||||
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
|
||||
app = ChatApp(temp_db_path, read_only=True)
|
||||
|
||||
async with app.run_test():
|
||||
chat_history = app.query_one(ChatHistory)
|
||||
|
||||
# Add a user message
|
||||
await chat_history.add_message("user", "Hello, how are you?")
|
||||
assert len(chat_history.messages) == 1
|
||||
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!")
|
||||
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
|
||||
async def test_chat_history_can_add_tool_calls(temp_db_path: Path):
|
||||
"""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
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
app, mock_client = _make_app(temp_db_path)
|
||||
|
||||
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
|
||||
app = ChatApp(temp_db_path, read_only=True)
|
||||
|
||||
async with app.run_test():
|
||||
chat_history = app.query_one(ChatHistory)
|
||||
|
||||
# Add a tool call
|
||||
tool_widget = await chat_history.add_tool_call("search", {"query": "test"})
|
||||
tool_widget = await chat_history.add_tool_call(
|
||||
"tool-1", "search", {"query": "test"}
|
||||
)
|
||||
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_widget)
|
||||
assert tool_widget._complete is True
|
||||
chat_history.mark_tool_complete("tool-1")
|
||||
assert tool_widget._completed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_history_can_add_citations(temp_db_path: Path):
|
||||
"""Test that ChatHistory can display inline citations."""
|
||||
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
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
app, mock_client = _make_app(temp_db_path)
|
||||
|
||||
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
|
||||
app = ChatApp(temp_db_path, read_only=True)
|
||||
|
||||
async with app.run_test():
|
||||
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)
|
||||
|
||||
# Verify citation widgets were added
|
||||
citation_widgets = chat_history.query(CitationWidget)
|
||||
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
|
||||
async def test_chat_history_thinking_indicator(temp_db_path: Path):
|
||||
"""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
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
app, mock_client = _make_app(temp_db_path)
|
||||
|
||||
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:
|
||||
chat_history = app.query_one(ChatHistory)
|
||||
|
||||
# Show thinking indicator
|
||||
await chat_history.show_thinking()
|
||||
thinking = chat_history.query(ThinkingWidget)
|
||||
assert len(list(thinking)) == 1
|
||||
|
||||
# Hide thinking indicator (remove is deferred, so pause)
|
||||
chat_history.hide_thinking()
|
||||
await pilot.pause()
|
||||
thinking = chat_history.query(ThinkingWidget)
|
||||
|
|
@ -247,67 +185,38 @@ async def test_chat_history_thinking_indicator(temp_db_path: Path):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_chat_resets_session(temp_db_path: Path):
|
||||
"""Test that clearing chat resets the session state."""
|
||||
from haiku.rag.chat.app import ChatApp
|
||||
async def test_clear_chat_resets_state(temp_db_path: Path):
|
||||
"""Test that clearing chat resets state and messages."""
|
||||
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)
|
||||
app, mock_client = _make_app(temp_db_path)
|
||||
|
||||
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:
|
||||
chat_history = app.query_one(ChatHistory)
|
||||
|
||||
# Add some messages
|
||||
await chat_history.add_message("user", "Hello")
|
||||
await chat_history.add_message("assistant", "Hi there")
|
||||
assert len(chat_history.messages) == 2
|
||||
|
||||
# Clear chat via action (available through command palette)
|
||||
await app.action_clear_chat()
|
||||
await pilot.pause()
|
||||
|
||||
# Verify messages cleared
|
||||
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
|
||||
async def test_citation_expand_collapse_with_enter(temp_db_path: Path):
|
||||
"""Test that pressing Enter on a focused citation toggles expand/collapse."""
|
||||
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
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
mock_client.__aexit__ = AsyncMock(return_value=None)
|
||||
app, mock_client = _make_app(temp_db_path)
|
||||
|
||||
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:
|
||||
chat_history = app.query_one(ChatHistory)
|
||||
|
||||
# Add a citation
|
||||
test_citation = Citation(
|
||||
index=1,
|
||||
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])
|
||||
|
||||
# Get the citation widget
|
||||
citation_widget = chat_history.query_one(CitationWidget)
|
||||
assert citation_widget.collapsed is True
|
||||
|
||||
# Focus the citation
|
||||
citation_widget.focus()
|
||||
await pilot.pause()
|
||||
|
||||
# Press Enter to expand
|
||||
await pilot.press("enter")
|
||||
await pilot.pause()
|
||||
assert citation_widget.collapsed is False
|
||||
|
||||
# Press Enter again to collapse
|
||||
await pilot.press("enter")
|
||||
await pilot.pause()
|
||||
assert citation_widget.collapsed is True
|
||||
|
|
|
|||
Loading…
Reference in a new issue