danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Clena citation handling now that we handle the total state

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-29 12:12:48 +02:00
parent 202df368a0
commit 77894b8c35
No known key found for this signature in database
2 changed files with 6 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -23,7 +23,6 @@ from haiku.rag.agents.chat.state import (
ChatDeps, ChatDeps,
ChatSessionState, ChatSessionState,
) )
from haiku.rag.agents.research.models import Citation
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
@ -103,7 +102,6 @@ class ChatApp(App):
self.session_state = ChatSessionState() 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._last_citations: list[Citation] = []
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] = []
@ -203,8 +201,6 @@ class ChatApp(App):
snapshot = getattr(meta_event, "snapshot", {}) snapshot = getattr(meta_event, "snapshot", {})
self._agui_state_snapshot = snapshot self._agui_state_snapshot = snapshot
chat_state = snapshot.get(AGUI_STATE_KEY, snapshot) chat_state = snapshot.get(AGUI_STATE_KEY, snapshot)
citations = chat_state.get("citations", [])
self._last_citations = [Citation(**c) for c in citations]
self._sync_session_state(chat_state) self._sync_session_state(chat_state)
elif meta_event.type == EventType.STATE_DELTA: elif meta_event.type == EventType.STATE_DELTA:
@ -217,8 +213,6 @@ class ChatApp(App):
chat_state = self._agui_state_snapshot.get( chat_state = self._agui_state_snapshot.get(
AGUI_STATE_KEY, self._agui_state_snapshot AGUI_STATE_KEY, self._agui_state_snapshot
) )
citations = chat_state.get("citations", [])
self._last_citations = [Citation(**c) for c in citations]
self._sync_session_state(chat_state) self._sync_session_state(chat_state)
async def _event_stream_handler( async def _event_stream_handler(
@ -253,7 +247,7 @@ class ChatApp(App):
# Clear for new query # Clear for new query
self._tool_call_widgets.clear() self._tool_call_widgets.clear()
self._last_citations.clear() self.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
@ -309,8 +303,8 @@ class ChatApp(App):
self._message_history = stream.all_messages() self._message_history = stream.all_messages()
# Add citations captured from tool metadata # Add citations captured from tool metadata
if self._last_citations: if self.session_state.citations:
await chat_history.add_citations(self._last_citations) await chat_history.add_citations(self.session_state.citations)
except asyncio.CancelledError: except asyncio.CancelledError:
chat_history.hide_thinking() chat_history.hide_thinking()
@ -329,7 +323,6 @@ class ChatApp(App):
"""Clear the chat history and reset session.""" """Clear the chat history and reset session."""
chat_history = self.query_one(ChatHistory) chat_history = self.query_one(ChatHistory)
await chat_history.clear_messages() await chat_history.clear_messages()
self._last_citations.clear()
self._message_history.clear() self._message_history.clear()
self._agui_state_snapshot = {} self._agui_state_snapshot = {}
# Reset context lock and session state (reset to CLI value) # Reset context lock and session state (reset to CLI value)

18
tests/chat/test_chat_app.py
View file

@ -337,10 +337,6 @@ async def test_handle_stream_event_extracts_citations_from_state_snapshot(
# Handle the event # Handle the event
await app._handle_stream_event(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"
# Session state should be synced # Session state should be synced
assert len(app.session_state.citations) == 1 assert len(app.session_state.citations) == 1
assert app.session_state.citations[0].chunk_id == "chunk1" assert app.session_state.citations[0].chunk_id == "chunk1"
@ -392,7 +388,7 @@ async def test_handle_stream_event_extracts_citations_from_state_delta(
) )
event1 = FunctionToolResultEvent(result=tool_return1) event1 = FunctionToolResultEvent(result=tool_return1)
await app._handle_stream_event(event1) await app._handle_stream_event(event1)
assert len(app._last_citations) == 0 assert len(app.session_state.citations) == 0
# Now handle a STATE_DELTA event that adds citations # Now handle a STATE_DELTA event that adds citations
delta_event = StateDeltaEvent( delta_event = StateDeltaEvent(
@ -430,14 +426,10 @@ async def test_handle_stream_event_extracts_citations_from_state_delta(
# Handle the delta event # Handle the delta event
await app._handle_stream_event(event2) 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"
# Session state should be synced # Session state should be synced
assert len(app.session_state.citations) == 1 assert len(app.session_state.citations) == 1
assert app.session_state.citations[0].chunk_id == "chunk1" assert app.session_state.citations[0].chunk_id == "chunk1"
assert app.session_state.citations[0].content == "Test content from delta"
@pytest.mark.asyncio @pytest.mark.asyncio
@ -514,14 +506,10 @@ async def test_handle_stream_event_delta_with_preinitialized_state(
# Handle the delta event # Handle the delta event
await app._handle_stream_event(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"
# Session state should be synced # Session state should be synced
assert len(app.session_state.citations) == 1 assert len(app.session_state.citations) == 1
assert app.session_state.citations[0].chunk_id == "chunk1" assert app.session_state.citations[0].chunk_id == "chunk1"
assert app.session_state.citations[0].content == "Content from first delta"
@pytest.mark.asyncio @pytest.mark.asyncio