Simplify TUI app
This commit is contained in:
parent
9e63ff1ff7
commit
0434cd068b
3 changed files with 55 additions and 441 deletions
|
|
@ -6,8 +6,6 @@ from datetime import datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
import jsonpatch
|
|
||||||
from ag_ui.core import EventType
|
|
||||||
from pydantic_ai import (
|
from pydantic_ai import (
|
||||||
Agent,
|
Agent,
|
||||||
AgentStreamEvent,
|
AgentStreamEvent,
|
||||||
|
|
@ -22,10 +20,6 @@ from haiku.rag.agents.chat.agent import (
|
||||||
create_chat_agent,
|
create_chat_agent,
|
||||||
trigger_background_summarization,
|
trigger_background_summarization,
|
||||||
)
|
)
|
||||||
from haiku.rag.agents.chat.state import (
|
|
||||||
AGUI_STATE_KEY,
|
|
||||||
ChatSessionState,
|
|
||||||
)
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config import get_config
|
from haiku.rag.config import get_config
|
||||||
from haiku.rag.tools.context import ToolContext
|
from haiku.rag.tools.context import ToolContext
|
||||||
|
|
@ -104,13 +98,11 @@ class ChatApp(App):
|
||||||
self.client: HaikuRAG | None = None
|
self.client: HaikuRAG | None = None
|
||||||
self.config = get_config()
|
self.config = get_config()
|
||||||
self.agent: Agent[ChatDeps, str] | None = None
|
self.agent: Agent[ChatDeps, str] | None = None
|
||||||
self.session_state = ChatSessionState()
|
|
||||||
self._is_processing = False
|
self._is_processing = False
|
||||||
self._tool_call_widgets: dict[str, Any] = {}
|
self._tool_call_widgets: dict[str, Any] = {}
|
||||||
self._current_worker: Worker[None] | None = None
|
self._current_worker: Worker[None] | None = None
|
||||||
self._message_history: list[ModelMessage] = []
|
self._message_history: list[ModelMessage] = []
|
||||||
self._document_filter: list[str] = []
|
self._document_filter: list[str] = []
|
||||||
self._agui_state_snapshot: dict[str, Any] = {}
|
|
||||||
|
|
||||||
def compose(self) -> "ComposeResult":
|
def compose(self) -> "ComposeResult":
|
||||||
"""Compose the UI layout."""
|
"""Compose the UI layout."""
|
||||||
|
|
@ -162,17 +154,11 @@ class ChatApp(App):
|
||||||
self.tool_context = ToolContext()
|
self.tool_context = ToolContext()
|
||||||
self.agent = create_chat_agent(self.config, self.client, self.tool_context)
|
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)
|
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
|
||||||
if session_state is not None:
|
if session_state is not None:
|
||||||
session_state.document_filter = self._document_filter
|
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
|
# Focus the input field
|
||||||
self.query_one(Input).focus()
|
self.query_one(Input).focus()
|
||||||
|
|
||||||
|
|
@ -181,40 +167,6 @@ class ChatApp(App):
|
||||||
if self.client:
|
if self.client:
|
||||||
await self.client.__aexit__(None, None, None)
|
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:
|
async def _handle_stream_event(self, event: AgentStreamEvent) -> None:
|
||||||
"""Handle streaming events from the agent."""
|
"""Handle streaming events from the agent."""
|
||||||
chat_history = self.query_one(ChatHistory)
|
chat_history = self.query_one(ChatHistory)
|
||||||
|
|
@ -232,32 +184,6 @@ class ChatApp(App):
|
||||||
widget = self._tool_call_widgets[tool_call_id]
|
widget = self._tool_call_widgets[tool_call_id]
|
||||||
chat_history.mark_tool_complete(widget)
|
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(
|
async def _event_stream_handler(
|
||||||
self,
|
self,
|
||||||
_ctx: RunContext[ChatDeps],
|
_ctx: RunContext[ChatDeps],
|
||||||
|
|
@ -290,7 +216,9 @@ class ChatApp(App):
|
||||||
|
|
||||||
# Clear for new query
|
# Clear for new query
|
||||||
self._tool_call_widgets.clear()
|
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
|
# Run agent in a worker to keep UI responsive
|
||||||
self._is_processing = True
|
self._is_processing = True
|
||||||
|
|
@ -310,39 +238,20 @@ class ChatApp(App):
|
||||||
await chat_history.show_thinking()
|
await chat_history.show_thinking()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Initialize AGUI state snapshot from session state for delta application
|
# Promote initial_context to QA session context on first run
|
||||||
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
|
|
||||||
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
|
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
|
||||||
|
|
||||||
qa_session_state = self.tool_context.get(
|
qa_session_state = self.tool_context.get(
|
||||||
QA_SESSION_NAMESPACE, QASessionState
|
QA_SESSION_NAMESPACE, QASessionState
|
||||||
)
|
)
|
||||||
if qa_session_state is not None:
|
if qa_session_state is not None:
|
||||||
if (
|
if not qa_session_state.session_context and self._initial_context:
|
||||||
not qa_session_state.session_context
|
qa_session_state.session_context = self._initial_context
|
||||||
and self.session_state.initial_context
|
|
||||||
):
|
|
||||||
qa_session_state.session_context = (
|
|
||||||
self.session_state.initial_context
|
|
||||||
)
|
|
||||||
|
|
||||||
deps = ChatDeps(
|
deps = ChatDeps(
|
||||||
config=self.config,
|
config=self.config,
|
||||||
tool_context=self.tool_context,
|
tool_context=self.tool_context,
|
||||||
is_new=False,
|
is_new=False,
|
||||||
state_key=AGUI_STATE_KEY,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
async with self.agent.run_stream(
|
async with self.agent.run_stream(
|
||||||
|
|
@ -367,27 +276,14 @@ class ChatApp(App):
|
||||||
# Update message history with this conversation
|
# Update message history with this conversation
|
||||||
self._message_history = stream.all_messages()
|
self._message_history = stream.all_messages()
|
||||||
|
|
||||||
# Add citations captured from tool metadata
|
# Add citations from ToolContext
|
||||||
if self.session_state.citations:
|
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
|
||||||
await chat_history.add_citations(self.session_state.citations)
|
if session_state and session_state.citations:
|
||||||
|
await chat_history.add_citations(session_state.citations)
|
||||||
|
|
||||||
# Trigger background summarization
|
# Trigger background summarization
|
||||||
trigger_background_summarization(deps)
|
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:
|
except asyncio.CancelledError:
|
||||||
chat_history.hide_thinking()
|
chat_history.hide_thinking()
|
||||||
await chat_history.add_message("assistant", "*Cancelled*")
|
await chat_history.add_message("assistant", "*Cancelled*")
|
||||||
|
|
@ -403,16 +299,18 @@ class ChatApp(App):
|
||||||
|
|
||||||
async def action_clear_chat(self) -> None:
|
async def action_clear_chat(self) -> None:
|
||||||
"""Clear the chat history and reset session."""
|
"""Clear the chat history and reset session."""
|
||||||
|
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
|
||||||
|
|
||||||
chat_history = self.query_one(ChatHistory)
|
chat_history = self.query_one(ChatHistory)
|
||||||
await chat_history.clear_messages()
|
await chat_history.clear_messages()
|
||||||
self._message_history.clear()
|
self._message_history.clear()
|
||||||
self._agui_state_snapshot = {}
|
|
||||||
# Reset context lock and session state (reset to CLI value)
|
|
||||||
self._context_locked = False
|
self._context_locked = False
|
||||||
self.session_state = ChatSessionState(
|
# Re-register fresh states in ToolContext
|
||||||
initial_context=self._initial_context,
|
self.tool_context.register(
|
||||||
document_filter=self._document_filter,
|
SESSION_NAMESPACE,
|
||||||
|
SessionState(document_filter=self._document_filter),
|
||||||
)
|
)
|
||||||
|
self.tool_context.register(QA_SESSION_NAMESPACE, QASessionState())
|
||||||
|
|
||||||
def action_focus_input(self) -> None:
|
def action_focus_input(self) -> None:
|
||||||
"""Focus the input field, or cancel if processing."""
|
"""Focus the input field, or cancel if processing."""
|
||||||
|
|
@ -462,29 +360,30 @@ class ChatApp(App):
|
||||||
|
|
||||||
async def action_show_context(self) -> None:
|
async def action_show_context(self) -> None:
|
||||||
"""Show context modal (edit initial context or view session context)."""
|
"""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.chat.widgets.context_modal import ContextModal
|
||||||
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
|
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)
|
qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState)
|
||||||
if qa_session_state is not None and qa_session_state.session_context:
|
if qa_session_state and qa_session_state.session_context:
|
||||||
from datetime import datetime
|
session_context = SessionContext(
|
||||||
|
|
||||||
from haiku.rag.agents.chat.state import SessionContext
|
|
||||||
|
|
||||||
self.session_state.session_context = SessionContext(
|
|
||||||
summary=qa_session_state.session_context,
|
summary=qa_session_state.session_context,
|
||||||
last_updated=datetime.now(),
|
last_updated=datetime.now(),
|
||||||
)
|
)
|
||||||
|
|
||||||
await self.push_screen(
|
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:
|
def on_context_modal_context_updated(self, event: Any) -> None:
|
||||||
"""Handle context updates from modal."""
|
"""Handle context updates from modal."""
|
||||||
if self.session_state and not self._context_locked:
|
if not self._context_locked:
|
||||||
self.session_state.initial_context = event.context or None
|
self._initial_context = event.context or None
|
||||||
|
|
||||||
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
|
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
|
||||||
"""Handle citation selection."""
|
"""Handle citation selection."""
|
||||||
|
|
@ -514,5 +413,6 @@ class ChatApp(App):
|
||||||
def on_document_filter_modal_filter_changed(self, event: Any) -> None:
|
def on_document_filter_modal_filter_changed(self, event: Any) -> None:
|
||||||
"""Handle document filter changes from modal."""
|
"""Handle document filter changes from modal."""
|
||||||
self._document_filter = event.selected
|
self._document_filter = event.selected
|
||||||
if self.session_state:
|
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
|
||||||
self.session_state.document_filter = self._document_filter
|
if session_state:
|
||||||
|
session_state.document_filter = self._document_filter
|
||||||
|
|
|
||||||
|
|
@ -8,7 +8,7 @@ from textual.screen import ModalScreen
|
||||||
from textual.widgets import Button, Markdown, Static, TextArea
|
from textual.widgets import Button, Markdown, Static, TextArea
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
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
|
class ContextModal(ModalScreen): # pragma: no cover
|
||||||
|
|
@ -80,20 +80,20 @@ class ContextModal(ModalScreen): # pragma: no cover
|
||||||
self.context = context
|
self.context = context
|
||||||
|
|
||||||
def __init__(
|
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:
|
) -> None:
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.session_state = session_state
|
self._initial_context = initial_context
|
||||||
|
self._session_context = session_context
|
||||||
self._is_locked = is_locked
|
self._is_locked = is_locked
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def _is_edit_mode(self) -> bool:
|
def _is_edit_mode(self) -> bool:
|
||||||
"""Edit mode when not locked and no session context yet."""
|
"""Edit mode when not locked and no session context yet."""
|
||||||
has_session_context = (
|
has_session_context = self._session_context and self._session_context.summary
|
||||||
self.session_state
|
|
||||||
and self.session_state.session_context
|
|
||||||
and self.session_state.session_context.summary
|
|
||||||
)
|
|
||||||
return not self._is_locked and not has_session_context
|
return not self._is_locked and not has_session_context
|
||||||
|
|
||||||
def compose(self) -> ComposeResult:
|
def compose(self) -> ComposeResult:
|
||||||
|
|
@ -105,9 +105,7 @@ class ContextModal(ModalScreen): # pragma: no cover
|
||||||
"This will be locked after you send your first message.",
|
"This will be locked after you send your first message.",
|
||||||
id="context-description",
|
id="context-description",
|
||||||
)
|
)
|
||||||
initial_value = ""
|
initial_value = self._initial_context or ""
|
||||||
if self.session_state and self.session_state.initial_context:
|
|
||||||
initial_value = self.session_state.initial_context
|
|
||||||
yield TextArea(initial_value, id="context-editor")
|
yield TextArea(initial_value, id="context-editor")
|
||||||
with Horizontal(id="button-row"):
|
with Horizontal(id="button-row"):
|
||||||
yield Button("Cancel", id="cancel-btn", variant="default")
|
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")
|
yield Button("Close", id="cancel-btn", variant="primary")
|
||||||
|
|
||||||
def _get_session_content(self) -> str:
|
def _get_session_content(self) -> str:
|
||||||
if not self.session_state:
|
if not self._session_context:
|
||||||
return "*No session state.*"
|
|
||||||
|
|
||||||
if not self.session_state.session_context:
|
|
||||||
return "*No session context yet. Ask a question first.*"
|
return "*No session context yet. Ask a question first.*"
|
||||||
|
|
||||||
ctx = self.session_state.session_context
|
ctx = self._session_context
|
||||||
updated = (
|
updated = (
|
||||||
ctx.last_updated.strftime("%Y-%m-%d %H:%M:%S")
|
ctx.last_updated.strftime("%Y-%m-%d %H:%M:%S")
|
||||||
if ctx.last_updated
|
if ctx.last_updated
|
||||||
|
|
|
||||||
|
|
@ -5,6 +5,7 @@ import pytest
|
||||||
from typer.testing import CliRunner
|
from typer.testing import CliRunner
|
||||||
|
|
||||||
from haiku.rag.cli import _cli as cli
|
from haiku.rag.cli import _cli as cli
|
||||||
|
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
|
||||||
|
|
||||||
runner = CliRunner()
|
runner = CliRunner()
|
||||||
|
|
||||||
|
|
@ -273,302 +274,20 @@ async def test_clear_chat_resets_session(temp_db_path: Path):
|
||||||
# Verify messages cleared
|
# Verify messages cleared
|
||||||
assert len(chat_history.messages) == 0
|
assert len(chat_history.messages) == 0
|
||||||
|
|
||||||
# Verify session state reset
|
# Verify ToolContext states are reset
|
||||||
assert app.session_state is not None
|
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
|
||||||
assert app.session_state.qa_history == []
|
|
||||||
assert app.session_state.citations == []
|
|
||||||
|
|
||||||
|
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
|
qa_session_state = app.tool_context.get(
|
||||||
async def test_handle_stream_event_extracts_citations_from_state_snapshot(
|
QA_SESSION_NAMESPACE, QASessionState
|
||||||
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."
|
|
||||||
)
|
)
|
||||||
|
assert qa_session_state is not None
|
||||||
|
assert qa_session_state.qa_history == []
|
||||||
|
assert qa_session_state.session_context is None
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue