Handle snapshot deltas in TUI
This commit is contained in:
Yiorgis Gozadinos
parent
3741274c6c
commit
e2158c7dad
2 changed files with 259 additions and 8 deletions
|
|
@ -6,6 +6,7 @@ from datetime import datetime
|
|||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import jsonpatch
|
||||
from ag_ui.core import EventType
|
||||
from pydantic_ai import (
|
||||
Agent,
|
||||
|
|
@ -106,6 +107,7 @@ class ChatApp(App):
|
|||
self._current_worker: Worker[None] | None = None
|
||||
self._message_history: list[ModelMessage] = []
|
||||
self._document_filter: list[str] = []
|
||||
self._agui_state_snapshot: dict[str, Any] = {}
|
||||
|
||||
def compose(self) -> "ComposeResult":
|
||||
"""Compose the UI layout."""
|
||||
|
|
@ -185,20 +187,33 @@ class ChatApp(App):
|
|||
widget = self._tool_call_widgets[tool_call_id]
|
||||
chat_history.mark_tool_complete(widget)
|
||||
|
||||
# Extract citations from StateSnapshotEvent in tool metadata
|
||||
# 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 (
|
||||
hasattr(meta_event, "type")
|
||||
and meta_event.type == EventType.STATE_SNAPSHOT
|
||||
):
|
||||
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._last_citations = [
|
||||
Citation(**c) for c in chat_state["citations"]
|
||||
]
|
||||
citations = chat_state.get("citations", [])
|
||||
self._last_citations = [Citation(**c) for c in citations]
|
||||
|
||||
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
|
||||
)
|
||||
citations = chat_state.get("citations", [])
|
||||
self._last_citations = [Citation(**c) for c in citations]
|
||||
|
||||
async def _event_stream_handler(
|
||||
self,
|
||||
|
|
@ -252,6 +267,12 @@ class ChatApp(App):
|
|||
await chat_history.show_thinking()
|
||||
|
||||
try:
|
||||
# Initialize AGUI state snapshot from session state for delta application
|
||||
if self.session_state:
|
||||
self._agui_state_snapshot = {
|
||||
AGUI_STATE_KEY: self.session_state.model_dump(mode="json")
|
||||
}
|
||||
|
||||
deps = ChatDeps(
|
||||
client=self.client,
|
||||
config=self.config,
|
||||
|
|
@ -304,6 +325,7 @@ class ChatApp(App):
|
|||
await chat_history.clear_messages()
|
||||
self._last_citations.clear()
|
||||
self._message_history.clear()
|
||||
self._agui_state_snapshot = {}
|
||||
# Reset context lock and session state (reset to CLI value)
|
||||
self._context_locked = False
|
||||
self.session_state = ChatSessionState(
|
||||
|
|
|
|||
|
|
@ -282,6 +282,235 @@ async def test_clear_chat_resets_session(temp_db_path: Path):
|
|||
assert app.session_state.session_id != original_session_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_stream_event_extracts_citations_from_state_snapshot(
|
||||
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: {
|
||||
"session_id": "test",
|
||||
"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)
|
||||
|
||||
# Citations should be extracted
|
||||
assert len(app._last_citations) == 1
|
||||
assert app._last_citations[0].chunk_id == "chunk1"
|
||||
|
||||
|
||||
@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: {
|
||||
"session_id": "test",
|
||||
"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._last_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)
|
||||
|
||||
# Citations should be extracted from the delta
|
||||
assert len(app._last_citations) == 1
|
||||
assert app._last_citations[0].chunk_id == "chunk1"
|
||||
assert app._last_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: {
|
||||
"session_id": "test",
|
||||
"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)
|
||||
|
||||
# Citations should be extracted from the delta
|
||||
assert len(app._last_citations) == 1
|
||||
assert app._last_citations[0].chunk_id == "chunk1"
|
||||
assert app._last_citations[0].content == "Content from first delta"
|
||||
|
||||
|
||||
@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."""
|
||||
|
|
|
|||
Loading…
Reference in a new issue