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

Update tests for state deltas

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-28 14:19:42 +02:00
parent e325962fcc
commit 639accc923
No known key found for this signature in database
1 changed files with 76 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

113
tests/agents/chat/test_chat_agent.py
View file

@ -1,6 +1,7 @@
from pathlib import Path
import pytest
from ag_ui.core import StateDeltaEvent, StateSnapshotEvent
from haiku.rag.agents.chat import (
AGUI_STATE_KEY,
@ -10,12 +11,29 @@ from haiku.rag.agents.chat import (
SearchAgent,
create_chat_agent,
)
from haiku.rag.agents.chat.context import get_cached_session_context
from haiku.rag.agents.chat.state import MAX_QA_HISTORY
from haiku.rag.agents.research.models import Citation
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
def extract_state_from_result(result, state_key: str = AGUI_STATE_KEY) -> dict | None:
"""Extract emitted state from agent result's tool return metadata."""
for message in result.all_messages():
if hasattr(message, "parts"):
for part in message.parts:
if hasattr(part, "metadata") and part.metadata:
for meta in part.metadata:
if isinstance(meta, StateSnapshotEvent):
return meta.snapshot.get(state_key)
elif isinstance(meta, StateDeltaEvent):
# For delta, we'd need to apply the patch
# For now, return None and let caller handle
pass
return None
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_chat_agent")
@ -449,11 +467,12 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
)
agent = create_chat_agent(Config)
session_state = ChatSessionState(session_id="test-citations")
# Pass session_state=None so agent creates fresh state with UUID
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
session_state=None,
state_key=AGUI_STATE_KEY,
)
# Ask a question that should use the ask tool with citations
@ -463,8 +482,12 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
)
assert result.output is not None
# Extract emitted state from result metadata
emitted_state = extract_state_from_result(result)
assert emitted_state is not None
# The qa_history should have been updated with the new Q&A
assert len(session_state.qa_history) >= 1
assert len(emitted_state.get("qa_history", [])) >= 1
@pytest.mark.asyncio
@ -483,16 +506,14 @@ async def test_chat_agent_ask_triggers_background_summarization(
)
agent = create_chat_agent(Config)
session_state = ChatSessionState(session_id="test-summarization")
# Pass session_state=None so agent creates fresh state with UUID
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
session_state=None,
state_key=AGUI_STATE_KEY,
)
# Initially no session_context
assert session_state.session_context is None
# Ask a question
result = await agent.run(
"What is the highest count class in the DocLayNet dataset?",
@ -500,19 +521,27 @@ async def test_chat_agent_ask_triggers_background_summarization(
)
assert result.output is not None
assert len(session_state.qa_history) >= 1
# Extract emitted state to get the session_id
emitted_state = extract_state_from_result(result)
assert emitted_state is not None
session_id = emitted_state.get("session_id")
assert session_id is not None
assert len(emitted_state.get("qa_history", [])) >= 1
# Wait for background task to complete
# The task should update session_state.session_context
# The task caches session_context server-side
cached_context = None
for _ in range(50): # Wait up to 5 seconds
if session_state.session_context is not None:
cached_context = get_cached_session_context(session_id)
if cached_context is not None:
break
await asyncio.sleep(0.1)
# Verify session_context was populated by background task
assert session_state.session_context is not None
assert session_state.session_context.summary != ""
assert session_state.session_context.last_updated is not None
assert cached_context is not None
assert cached_context.summary != ""
assert cached_context.last_updated is not None
@pytest.mark.asyncio
@ -536,52 +565,62 @@ async def test_chat_agent_ask_with_prior_answer_retrieval(
)
agent = create_chat_agent(Config)
session_state = ChatSessionState(session_id="test-prior-answers")
deps = ChatDeps(
# First call with no session state
deps1 = ChatDeps(
client=client,
config=Config,
session_state=session_state,
session_state=None,
state_key=AGUI_STATE_KEY,
)
# First ask - establishes qa_history
result1 = await agent.run(
"What are the class labels in DocLayNet?",
deps=deps,
deps=deps1,
)
assert result1.output is not None
assert len(session_state.qa_history) == 1
# First question should NOT have embedding yet (set lazily on next ask)
assert session_state.qa_history[0].question_embedding is None
# Extract emitted state from first call
state1 = extract_state_from_result(result1)
assert state1 is not None
assert len(state1.get("qa_history", [])) == 1
session_id = state1.get("session_id")
assert session_id is not None
# Wait for background summarization to complete
for _ in range(50):
if session_state.session_context is not None:
if get_cached_session_context(session_id) is not None:
break
await asyncio.sleep(0.1)
# Create new session state from emitted state for second call
# (simulating client sending state back to server)
session_state2 = ChatSessionState(
session_id=session_id,
qa_history=[QAResponse(**qa) for qa in state1.get("qa_history", [])],
citation_registry=state1.get("citation_registry", {}),
)
deps2 = ChatDeps(
client=client,
config=Config,
session_state=session_state2,
state_key=AGUI_STATE_KEY,
)
# Second ask - similar question triggers prior answer retrieval
# This will embed the first question and compare similarity
result2 = await agent.run(
"Tell me about DocLayNet class labels",
deps=deps,
deps=deps2,
)
assert result2.output is not None
# qa_history should now have 2 entries
assert len(session_state.qa_history) == 2
# First question should now have embedding (set during second ask's recall check)
assert session_state.qa_history[0].question_embedding is not None
# The important thing is that prior answer retrieval happened
# We can verify this by checking session_state2 was used (embedding added)
assert session_state2.qa_history[0].question_embedding is not None
# Embedding should be a list of floats
assert isinstance(session_state.qa_history[0].question_embedding, list)
assert len(session_state.qa_history[0].question_embedding) > 0
# Verify prior answer was reused without new searches:
# Second answer's citations should be subset of first answer's citations
first_chunk_ids = {c.chunk_id for c in session_state.qa_history[0].citations}
second_chunk_ids = {c.chunk_id for c in session_state.qa_history[1].citations}
assert second_chunk_ids <= first_chunk_ids, (
"Second answer should reuse prior citations, not perform new searches"
)
assert isinstance(session_state2.qa_history[0].question_embedding, list)
assert len(session_state2.qa_history[0].question_embedding) > 0
def test_fifo_limit_enforcement():