danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
haiku.rag/tests/agents/chat/test_chat_agent.py
Yiorgis Gozadinos a870052ee9
Fix coverage
2026-02-17 13:22:23 +02:00

1627 lines
54 KiB
Python
Raw Blame History

from pathlib import Path
import pytest
from ag_ui.core import StateDeltaEvent
from haiku.rag.agents.chat import (
AGUI_STATE_KEY,
ChatDeps,
ChatSessionState,
create_chat_agent,
prepare_chat_context,
run_chat_agent,
)
from haiku.rag.agents.chat.context import _summarization_tasks
from haiku.rag.agents.research.models import Citation
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.tools import ToolContext
from haiku.rag.tools.qa import MAX_QA_HISTORY, QAHistoryEntry
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionContext, SessionState
def extract_state_from_result(result, state_key: str = AGUI_STATE_KEY) -> dict | None:
"""Extract emitted state from agent result's tool return metadata.
Applies the JSON Patch delta to an empty state to get the final state.
"""
import jsonpatch
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, StateDeltaEvent):
empty_state = {
state_key: ChatSessionState().model_dump(mode="json")
}
patched = jsonpatch.apply_patch(empty_state, meta.delta)
return patched.get(state_key)
return None
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_chat_agent")
def test_create_chat_agent(temp_db_path):
"""Test that create_chat_agent returns a properly configured agent."""
agent = create_chat_agent(Config)
assert agent is not None
assert agent.name == "chat_agent" or agent.name is None
def test_chat_deps_initialization(temp_db_path):
"""Test ChatDeps can be initialized with required fields."""
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
deps = ChatDeps(config=Config, client=client, tool_context=context)
assert deps.config is Config
assert deps.client is client
assert deps.tool_context is context
client.close()
def test_chat_deps_is_agent_deps(temp_db_path):
"""Test ChatDeps is a subclass of AgentDeps."""
from haiku.rag.tools.deps import AgentDeps
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
deps = ChatDeps(config=Config, client=client, tool_context=context)
assert isinstance(deps, AgentDeps)
client.close()
def test_agui_state_key_constant():
"""Test AGUI_STATE_KEY is exported with correct value."""
assert AGUI_STATE_KEY == "haiku.rag.chat"
def test_chat_deps_state_setter_none(temp_db_path):
"""Test ChatDeps.state setter handles None gracefully."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
context.register(QA_SESSION_NAMESPACE, QASessionState())
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(config=Config, client=client, tool_context=context)
# Setting state to None should be a no-op
deps.state = None
# State should remain unchanged
qa_session_state = context.get(QA_SESSION_NAMESPACE)
assert isinstance(qa_session_state, QASessionState)
assert qa_session_state.session_context is None
client.close()
def test_chat_deps_state_setter_handles_initial_context(temp_db_path):
"""Test ChatDeps.state setter transfers initial_context to qa_session_state."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
context.state_key = AGUI_STATE_KEY
# Register QASessionState (normally done by prepare_chat_context)
context.register(QA_SESSION_NAMESPACE, QASessionState())
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(config=Config, client=client, tool_context=context)
# Client sends initial_context with no session_context
incoming_state = {
AGUI_STATE_KEY: {
"initial_context": "Background info about the project",
"session_context": None,
"qa_history": [],
"citations": [],
"document_filter": [],
"citation_registry": {},
}
}
deps.state = incoming_state
# initial_context should be copied to qa_session_state.session_context
qa_session_state = context.get(QA_SESSION_NAMESPACE)
assert isinstance(qa_session_state, QASessionState)
assert qa_session_state.session_context is not None
assert (
qa_session_state.session_context.summary == "Background info about the project"
)
client.close()
def test_chat_deps_state_setter_parses_session_context_dict(temp_db_path):
"""Test ChatDeps.state setter parses session_context dict and extracts summary."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
context.state_key = AGUI_STATE_KEY
context.register(QA_SESSION_NAMESPACE, QASessionState())
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(config=Config, client=client, tool_context=context)
# Client sends session_context as a dict (as it comes from JSON)
incoming_state = {
AGUI_STATE_KEY: {
"session_context": {
"summary": "Previous conversation summary",
"last_updated": "2025-01-27T12:00:00",
},
"qa_history": [],
"citations": [],
"document_filter": [],
"citation_registry": {},
}
}
deps.state = incoming_state
# session_context dict should be parsed into SessionContext
qa_session_state = context.get(QA_SESSION_NAMESPACE)
assert isinstance(qa_session_state, QASessionState)
assert qa_session_state.session_context is not None
assert isinstance(qa_session_state.session_context, SessionContext)
assert qa_session_state.session_context.summary == "Previous conversation summary"
client.close()
def test_chat_deps_state_setter_preserves_server_session_context(temp_db_path):
"""Test that server's session_context is preferred over client's stale value."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
context.state_key = AGUI_STATE_KEY
qa_state = QASessionState()
qa_state.session_context = SessionContext(
summary="Fresh summary from background summarizer"
)
context.register(QA_SESSION_NAMESPACE, qa_state)
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(config=Config, client=client, tool_context=context)
# Client sends stale session_context
incoming_state = {
AGUI_STATE_KEY: {
"session_context": {
"summary": "Stale summary from client",
"last_updated": "2025-01-27T12:00:00",
},
"qa_history": [],
"citations": [],
"document_filter": [],
"citation_registry": {},
}
}
deps.state = incoming_state
# Server's fresher session_context should be preserved
qa_session_state = context.get(QA_SESSION_NAMESPACE)
assert isinstance(qa_session_state, QASessionState)
assert qa_session_state.session_context is not None
assert (
qa_session_state.session_context.summary
== "Fresh summary from background summarizer"
)
client.close()
def test_trigger_background_summarization_no_qa_history(temp_db_path):
"""Test trigger_background_summarization returns early when qa_history is empty."""
from haiku.rag.agents.chat.agent import trigger_background_summarization
from haiku.rag.agents.chat.context import _summarization_tasks
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
context.register(QA_SESSION_NAMESPACE, QASessionState())
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(config=Config, client=client, tool_context=context)
tasks_before = len(_summarization_tasks)
trigger_background_summarization(deps)
assert len(_summarization_tasks) == tasks_before
client.close()
def test_chat_session_state():
"""Test ChatSessionState model."""
state = ChatSessionState()
assert state.citations == []
assert state.qa_history == []
def test_citation():
"""Test Citation model."""
citation = Citation(
index=1,
document_id="doc-123",
chunk_id="chunk-456",
document_uri="test.md",
document_title="Test Document",
page_numbers=[1, 2],
headings=["Section 1"],
content="Test content",
)
assert citation.index == 1
assert citation.document_id == "doc-123"
assert citation.chunk_id == "chunk-456"
assert citation.content == "Test content"
def test_qa_response():
"""Test QAHistoryEntry model."""
citation = Citation(
index=1,
document_id="doc-123",
chunk_id="chunk-456",
document_uri="test.md",
document_title="Test Document",
content="Test content",
)
qa = QAHistoryEntry(
question="What is this?",
answer="This is a test",
confidence=0.95,
citations=[citation],
)
assert qa.question == "What is this?"
assert qa.answer == "This is a test"
assert qa.confidence == 0.95
assert len(qa.citations) == 1
assert qa.sources == ["Test Document"]
def test_qa_response_sources_with_uri_fallback():
"""Test QAHistoryEntry.sources falls back to URI when title is None."""
citation = Citation(
index=1,
document_id="doc-123",
chunk_id="chunk-456",
document_uri="test.md",
document_title=None,
content="Test content",
)
qa = QAHistoryEntry(
question="What is this?",
answer="This is a test",
citations=[citation],
)
assert qa.sources == ["test.md"]
def test_qa_response_to_search_answer():
"""Test QAHistoryEntry.to_search_answer() converts to SearchAnswer for research graph."""
citation = Citation(
index=1,
document_id="doc-123",
chunk_id="chunk-456",
document_uri="test.md",
document_title="Test Document",
content="Test content",
)
qa = QAHistoryEntry(
question="What is the answer?",
answer="The answer is 42",
confidence=0.95,
citations=[citation],
)
search_answer = qa.to_search_answer()
assert search_answer.query == "What is the answer?"
assert search_answer.answer == "The answer is 42"
assert search_answer.confidence == 0.95
assert search_answer.cited_chunks == ["chunk-456"]
assert len(search_answer.citations) == 1
assert search_answer.citations[0].chunk_id == "chunk-456"
# DocLayNet content for testing
DOCLAYNET_CLASS_LABELS = """
DocLayNet Dataset - Class Labels
DocLayNet defines 11 distinct class labels for document layout analysis:
1. Caption - Text describing figures or tables
2. Footnote - Notes at the bottom of pages
3. Formula - Mathematical expressions
4. List-item - Items in bulleted or numbered lists
5. Page-footer - Footer content on pages
6. Page-header - Header content on pages
7. Picture - Images and diagrams
8. Section-header - Headings for document sections
9. Table - Tabular data
10. Text - Regular paragraph text (highest count: 510,377 instances)
11. Title - Document titles
The Text class has the highest count with 510,377 instances in the dataset.
"""
DOCLAYNET_ANNOTATION = """
DocLayNet Dataset - Annotation Process
The annotation process was organized into 4 phases:
- Phase 1: Data selection and preparation by a small team of experts
- Phase 2: Label selection and guideline definition
- Phase 3: Annotation by 40 dedicated annotators
- Phase 4: Quality control and continuous supervision
The Corpus Conversion Service (CCS) was used for annotation, providing a visual interface.
"""
DOCLAYNET_DATA_SOURCES = """
DocLayNet Dataset - Data Sources
The data sources for DocLayNet include:
- Publication repositories such as arXiv
- Government offices and official documents
- Company websites and corporate reports
- Data directory services for financial reports
- Patent documents
Scanned documents were excluded to avoid rotation and skewing issues.
"""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_run_chat_agent(allow_model_requests, temp_db_path):
"""Test run_chat_agent returns agent output string."""
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
output = await run_chat_agent(agent, deps, "Search for class labels")
assert isinstance(output, str)
assert len(output) > 0
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_search_tool(allow_model_requests, temp_db_path):
"""Test the chat agent's search tool functionality."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_ANNOTATION,
uri="doclaynet-annotation",
title="DocLayNet Annotation",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Ask something that should trigger the search tool
result = await agent.run(
"Search for documents about class labels",
deps=deps,
)
assert result.output is not None
assert len(result.output) > 0
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_search_tool_with_filter(allow_model_requests, temp_db_path):
"""Test the chat agent's search tool with document filter."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_DATA_SOURCES,
uri="doclaynet-sources",
title="DocLayNet Sources",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Ask to search within a specific document
result = await agent.run(
"Search for information about class labels in the DocLayNet Class Labels document",
deps=deps,
)
assert result.output is not None
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_get_document_tool(allow_model_requests, temp_db_path):
"""Test the chat agent's get_document tool."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add a test document
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Ask to get a specific document
result = await agent.run(
"Get me the DocLayNet Class Labels document",
deps=deps,
)
assert result.output is not None
# The response should contain info about the document
assert "DocLayNet" in result.output or "class" in result.output.lower()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_get_document_not_found(allow_model_requests, temp_db_path):
"""Test the chat agent's get_document tool when document is not found."""
async with HaikuRAG(temp_db_path, create=True) as client:
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Ask for a document that doesn't exist
result = await agent.run(
"Get me the nonexistent document",
deps=deps,
)
assert result.output is not None
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path):
"""Test that the ask tool is called and can add citations to session state."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add a document with specific content
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Ask a question that should use the ask tool
result = await agent.run(
"What is the highest count class in the DocLayNet dataset?",
deps=deps,
)
assert result.output is not None
# Verify the agent used the ask tool by checking for tool calls
tool_calls = [
part
for msg in result.all_messages()
if hasattr(msg, "parts")
for part in msg.parts
if hasattr(part, "tool_name") and part.tool_name == "ask"
]
assert len(tool_calls) >= 1, "Expected ask tool to be called"
# Session state should be registered (citations may or may not be present
# depending on whether the research graph found relevant evidence)
session_state = context.get(SESSION_NAMESPACE)
assert isinstance(session_state, SessionState)
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_ask_triggers_background_summarization(
allow_model_requests, temp_db_path
):
"""Test that the ask tool triggers background session context summarization.
Patches the internal trigger in run_qa_core to avoid concurrent HTTP calls
that break VCR cassette replay ordering. Triggers summarization explicitly
after the agent run completes.
"""
from unittest.mock import patch
from haiku.rag.agents.chat.agent import trigger_background_summarization
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Patch internal trigger to avoid concurrent HTTP calls during VCR
with patch("haiku.rag.agents.chat.agent._trigger_summarization"):
result = await agent.run(
"What is the highest count class in the DocLayNet dataset?",
deps=deps,
)
assert result.output is not None
# Trigger summarization explicitly (sequential, deterministic)
trigger_background_summarization(deps)
# Wait for background task to complete
qa_session_state = context.get(QA_SESSION_NAMESPACE, QASessionState)
assert qa_session_state is not None
key = id(qa_session_state)
if key in _summarization_tasks:
await _summarization_tasks[key]
# Verify session_context was populated by background task
assert qa_session_state.session_context is not None
assert qa_session_state.session_context.summary != ""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_path):
"""Test multi-turn conversation with initial context, summarization, and prior recall.
Exercises the full conversation flow:
1. Initial context is transferred to session context
2. First question triggers background summarization
3. Second related question uses prior answer recall and updated session context
4. Both qa_history entries are present after two turns
The ask tool internally fires background summarization (concurrent HTTP calls)
which causes VCR cassette mismatches. We patch it to a no-op and trigger
summarization explicitly after each turn to keep HTTP ordering deterministic.
"""
from unittest.mock import patch
from haiku.rag.agents.chat.agent import trigger_background_summarization
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_ANNOTATION,
uri="doclaynet-annotation",
title="DocLayNet Annotation",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Set initial state with initial_context (mimicking AG-UI client)
deps.state = {
AGUI_STATE_KEY: {
"initial_context": "The user is researching the DocLayNet dataset for a paper on document layout analysis.",
"session_context": None,
"qa_history": [],
"citations": [],
"document_filter": [],
"citation_registry": {},
}
}
# initial_context should be transferred to QASessionState
qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState)
assert qa_session is not None
assert qa_session.session_context is not None
assert (
qa_session.session_context.summary
== "The user is researching the DocLayNet dataset for a paper on document layout analysis."
)
# Patch the internal summarization trigger in the ask tool to avoid
# concurrent HTTP calls that break VCR cassette replay ordering.
with patch(
"haiku.rag.agents.chat.agent._trigger_summarization",
):
# First question about class labels
result1 = await agent.run(
"What are the class labels defined in DocLayNet?",
deps=deps,
)
assert result1.output is not None
# Trigger summarization explicitly (sequential, no concurrency)
trigger_background_summarization(deps)
key = id(qa_session)
if key in _summarization_tasks:
await _summarization_tasks[key]
assert qa_session.session_context is not None
assert qa_session.session_context.summary != ""
# qa_history should have one entry
qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState)
assert qa_session is not None
assert len(qa_session.qa_history) >= 1
# Second related question - uses prior answers and updated session context
with patch(
"haiku.rag.agents.chat.agent._trigger_summarization",
):
result2 = await agent.run(
"How were the annotations created and how many annotators were involved?",
deps=deps,
message_history=result1.all_messages(),
)
assert result2.output is not None
# Trigger summarization explicitly
trigger_background_summarization(deps)
qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState)
assert qa_session is not None
key = id(qa_session)
if key in _summarization_tasks:
await _summarization_tasks[key]
# qa_history should have two entries
assert len(qa_session.qa_history) >= 2
# Session context should be updated with newer summary
assert qa_session.session_context is not None
assert qa_session.session_context.summary != ""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_ask_with_prior_answer_retrieval(
allow_model_requests, temp_db_path
):
"""Test that ask tool retrieves relevant prior answers from qa_history.
This exercises the prior answer retrieval logic:
1. First ask populates qa_history with question_embedding
2. Second similar ask should find the prior answer via embedding similarity
"""
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps1 = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# First ask - establishes qa_history
result1 = await agent.run(
"What are the class labels in DocLayNet?",
deps=deps1,
)
assert result1.output is not None
# Check that session state has citations after first call
session_state = context.get(SESSION_NAMESPACE)
assert isinstance(session_state, SessionState)
# Citations might be 0 if the answer came from prior context
assert len(session_state.citations) >= 0
# Second ask - similar question triggers prior answer retrieval
result2 = await agent.run(
"Tell me about DocLayNet class labels",
deps=deps1,
)
assert result2.output is not None
# The QA session state should have history entries
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
qa_session = context.get(QA_SESSION_NAMESPACE)
assert isinstance(qa_session, QASessionState)
# After two asks, we should have entries in qa_history
assert len(qa_session.qa_history) >= 1
def test_fifo_limit_enforcement():
"""Test that FIFO limit enforcement logic works correctly.
This tests the FIFO trimming logic used in the ask() tool:
if len(qa_history) > MAX_QA_HISTORY:
qa_history = qa_history[-MAX_QA_HISTORY:]
"""
# Create a session state with MAX_QA_HISTORY + 1 entries
qa_history = [
QAHistoryEntry(
question=f"Question {i}",
answer=f"Answer {i}",
confidence=0.9,
)
for i in range(MAX_QA_HISTORY + 1)
]
session_state = ChatSessionState(
qa_history=qa_history,
)
# Simulate the FIFO enforcement from agent.py
if len(session_state.qa_history) > MAX_QA_HISTORY:
session_state.qa_history = session_state.qa_history[-MAX_QA_HISTORY:]
# History should be trimmed to MAX_QA_HISTORY
assert len(session_state.qa_history) == MAX_QA_HISTORY
# The first entry should now be "Question 1" (Question 0 was dropped)
assert session_state.qa_history[0].question == "Question 1"
# The last entry should be the last added question
assert session_state.qa_history[-1].question == f"Question {MAX_QA_HISTORY}"
def test_chat_session_state_document_filter():
"""Test ChatSessionState with document_filter."""
state = ChatSessionState(
document_filter=["doc1.pdf", "doc2.pdf"],
)
assert state.document_filter == ["doc1.pdf", "doc2.pdf"]
def test_chat_session_state_document_filter_default_empty():
"""Test ChatSessionState document_filter defaults to empty list."""
state = ChatSessionState()
assert state.document_filter == []
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_search_with_session_filter(
allow_model_requests, temp_db_path
):
"""Test that session document_filter restricts search results."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add two distinct documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_DATA_SOURCES,
uri="doclaynet-sources",
title="DocLayNet Sources",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
# Set session filter to only include the labels document
session_state = context.get(SESSION_NAMESPACE)
assert isinstance(session_state, SessionState)
session_state.document_filter = ["DocLayNet Class Labels"]
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Search should only return results from the filtered document
result = await agent.run(
"Search for information about DocLayNet",
deps=deps,
)
assert result.output is not None
# Check that citations in context are only from the filtered document
session_state = context.get(SESSION_NAMESPACE)
assert isinstance(session_state, SessionState)
# If citations were added, they should only be from the labels document
for citation in session_state.citations:
assert "labels" in citation.document_uri.lower() or "Labels" in (
citation.document_title or ""
)
def test_ask_tool_citation_registry_logic():
"""Test the citation index assignment logic used by the ask tool.
Verifies that:
1. First chunk gets index 1
2. Second unique chunk gets index 2
3. Same chunk_id always gets same index
4. Indices don't reset between calls
"""
session_state = SessionState()
# Simulate first ask tool building citations
first_ask_chunks = ["chunk-a", "chunk-b"]
first_citations = []
for chunk_id in first_ask_chunks:
index = session_state.get_or_assign_index(chunk_id)
first_citations.append(
Citation(
index=index,
document_id="doc-1",
chunk_id=chunk_id,
document_uri="test.md",
content="test",
)
)
assert first_citations[0].index == 1
assert first_citations[1].index == 2
# Simulate second ask tool - overlapping chunk_id should keep same index
second_ask_chunks = ["chunk-b", "chunk-c"] # chunk-b was in first ask
second_citations = []
for chunk_id in second_ask_chunks:
index = session_state.get_or_assign_index(chunk_id)
second_citations.append(
Citation(
index=index,
document_id="doc-1",
chunk_id=chunk_id,
document_uri="test.md",
content="test",
)
)
# chunk-b should have same index as before
assert second_citations[0].index == 2
# chunk-c is new, gets next index
assert second_citations[1].index == 3
# Registry should have all three chunks
assert len(session_state.citation_registry) == 3
assert session_state.citation_registry == {"chunk-a": 1, "chunk-b": 2, "chunk-c": 3}
def test_search_tool_citation_registry_logic():
"""Test the citation index assignment logic used by the search tool.
Verifies that search and ask tools share the same registry,
maintaining stable indices across different tool calls.
"""
session_state = SessionState()
# Simulate ask tool first (assigns indices 1, 2)
for chunk_id in ["chunk-a", "chunk-b"]:
session_state.get_or_assign_index(chunk_id)
# Simulate search tool returning overlapping + new chunks
search_chunks = ["chunk-b", "chunk-c", "chunk-d"] # chunk-b already exists
search_citations = []
for chunk_id in search_chunks:
index = session_state.get_or_assign_index(chunk_id)
search_citations.append(
Citation(
index=index,
document_id="doc-1",
chunk_id=chunk_id,
document_uri="test.md",
content="test",
)
)
# chunk-b should have same index as assigned by ask (2)
assert search_citations[0].index == 2
# New chunks get incrementing indices
assert search_citations[1].index == 3
assert search_citations[2].index == 4
# Registry should have all four chunks
assert len(session_state.citation_registry) == 4
assert session_state.citation_registry == {
"chunk-a": 1,
"chunk-b": 2,
"chunk-c": 3,
"chunk-d": 4,
}
# =============================================================================
# Prior Answer Recall Tests
# =============================================================================
def test_cosine_similarity_identical_vectors():
"""Test cosine similarity returns 1.0 for identical vectors."""
from haiku.rag.tools.qa import _cosine_similarity
vec = [1.0, 2.0, 3.0]
assert _cosine_similarity(vec, vec) == pytest.approx(1.0)
def test_cosine_similarity_orthogonal_vectors():
"""Test cosine similarity returns 0.0 for orthogonal vectors."""
from haiku.rag.tools.qa import _cosine_similarity
vec1 = [1.0, 0.0, 0.0]
vec2 = [0.0, 1.0, 0.0]
assert _cosine_similarity(vec1, vec2) == pytest.approx(0.0)
def test_cosine_similarity_opposite_vectors():
"""Test cosine similarity returns -1.0 for opposite vectors."""
from haiku.rag.tools.qa import _cosine_similarity
vec1 = [1.0, 2.0, 3.0]
vec2 = [-1.0, -2.0, -3.0]
assert _cosine_similarity(vec1, vec2) == pytest.approx(-1.0)
def test_cosine_similarity_zero_vector():
"""Test cosine similarity handles zero vectors gracefully."""
from haiku.rag.tools.qa import _cosine_similarity
vec = [1.0, 2.0, 3.0]
zero = [0.0, 0.0, 0.0]
assert _cosine_similarity(vec, zero) == 0.0
assert _cosine_similarity(zero, vec) == 0.0
assert _cosine_similarity(zero, zero) == 0.0
def test_prior_answer_relevance_threshold_constant():
"""Test PRIOR_ANSWER_RELEVANCE_THRESHOLD is set to expected value."""
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD
assert PRIOR_ANSWER_RELEVANCE_THRESHOLD == 0.7
def test_prior_answer_matching_above_threshold():
"""Test that similar questions (above threshold) are matched."""
from haiku.rag.tools.qa import (
PRIOR_ANSWER_RELEVANCE_THRESHOLD,
_cosine_similarity,
)
# Simulate two nearly identical question embeddings
question_embedding = [0.5, 0.5, 0.5, 0.5]
prior_embedding = [0.51, 0.49, 0.5, 0.5] # Very similar
similarity = _cosine_similarity(question_embedding, prior_embedding)
assert similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD
def test_prior_answer_matching_below_threshold():
"""Test that dissimilar questions (below threshold) are not matched."""
from haiku.rag.tools.qa import (
PRIOR_ANSWER_RELEVANCE_THRESHOLD,
_cosine_similarity,
)
# Simulate two different question embeddings
question_embedding = [1.0, 0.0, 0.0, 0.0]
prior_embedding = [0.0, 1.0, 0.0, 0.0] # Orthogonal = very different
similarity = _cosine_similarity(question_embedding, prior_embedding)
assert similarity < PRIOR_ANSWER_RELEVANCE_THRESHOLD
def test_qa_response_embedding_cache():
"""Test that QAHistoryEntry stores and retrieves question_embedding correctly."""
embedding = [0.1, 0.2, 0.3, 0.4]
qa = QAHistoryEntry(
question="What is X?",
answer="X is Y.",
confidence=0.9,
question_embedding=embedding,
)
assert qa.question_embedding == embedding
# Embedding should be excluded from serialization (AG-UI state)
serialized = qa.model_dump()
assert "question_embedding" not in serialized
def test_qa_response_embedding_default_none():
"""Test that QAHistoryEntry.question_embedding defaults to None."""
qa = QAHistoryEntry(
question="What is X?",
answer="X is Y.",
confidence=0.9,
)
assert qa.question_embedding is None
# =============================================================================
# Background Task Cancellation Tests
# =============================================================================
@pytest.mark.asyncio
async def test_summarization_task_cancellation():
"""Test that new summarization tasks cancel previous ones for same state object."""
import asyncio
from haiku.rag.agents.chat.context import _summarization_tasks
# Clear any existing tasks
_summarization_tasks.clear()
key = 12345 # Simulates id(qa_session_state)
# Create a slow task that simulates summarization
async def slow_task():
await asyncio.sleep(10) # Would take 10 seconds
# Start first task
task1 = asyncio.create_task(slow_task())
_summarization_tasks[key] = task1
# Simulate what happens when second ask comes in - cancel first task
if key in _summarization_tasks:
_summarization_tasks[key].cancel()
# Yield to let cancellation propagate
await asyncio.sleep(0)
# Start second task
task2 = asyncio.create_task(slow_task())
_summarization_tasks[key] = task2
# First task should be cancelled
assert task1.cancelled() or task1.done()
# Second task should be running
assert not task2.done()
# Cleanup
task2.cancel()
try:
await task2
except asyncio.CancelledError:
pass
_summarization_tasks.clear()
# =============================================================================
# list_documents Tool Tests
# =============================================================================
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_list_documents_basic(allow_model_requests, temp_db_path):
"""Test that list_documents tool returns available documents."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_ANNOTATION,
uri="doclaynet-annotation",
title="DocLayNet Annotation",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Ask to list documents
result = await agent.run(
"What documents are available in the knowledge base?",
deps=deps,
)
assert result.output is not None
# Should mention both documents
assert "DocLayNet" in result.output
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_list_documents_with_session_filter(allow_model_requests, temp_db_path):
"""Test that list_documents respects session document_filter."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_DATA_SOURCES,
uri="doclaynet-sources",
title="DocLayNet Sources",
)
# Set session filter to only include the labels document
context = ToolContext()
context.register(
SESSION_NAMESPACE,
SessionState(document_filter=["DocLayNet Class Labels"]),
)
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Ask to list documents - should only show filtered documents
result = await agent.run(
"Show me what documents are available",
deps=deps,
)
assert result.output is not None
# Should only mention the Labels document, not Sources
assert "Labels" in result.output or "labels" in result.output
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_list_documents_pagination(allow_model_requests, temp_db_path):
"""Test that list_documents supports pagination via limit/offset."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add multiple test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_ANNOTATION,
uri="doclaynet-annotation",
title="DocLayNet Annotation",
)
await client.create_document(
content=DOCLAYNET_DATA_SOURCES,
uri="doclaynet-sources",
title="DocLayNet Sources",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Ask to list first 2 documents
result = await agent.run(
"List the first 2 documents available",
deps=deps,
)
assert result.output is not None
# =============================================================================
# summarize_document Tool Tests
# =============================================================================
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_summarize_document_found(allow_model_requests, temp_db_path):
"""Test that summarize_document generates a summary for a found document."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add a test document
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Ask to summarize a specific document
result = await agent.run(
"Summarize the DocLayNet Class Labels document",
deps=deps,
)
assert result.output is not None
# Should contain summary content about class labels
assert len(result.output) > 50
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_summarize_document_not_found(allow_model_requests, temp_db_path):
"""Test that summarize_document handles not found documents gracefully."""
async with HaikuRAG(temp_db_path, create=True) as client:
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
deps = ChatDeps(
config=Config,
client=client,
tool_context=context,
)
# Ask to summarize a document that doesn't exist
result = await agent.run(
"Summarize the nonexistent document",
deps=deps,
)
assert result.output is not None
# Should indicate the document wasn't found
# =============================================================================
# count_documents Tests
# =============================================================================
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_count_documents(temp_db_path):
"""Test count_documents method."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Empty database
assert await client.count_documents() == 0
# Add documents
await client.create_document(content="Doc 1", uri="test/doc1.pdf")
await client.create_document(content="Doc 2", uri="test/doc2.pdf")
await client.create_document(content="Doc 3", uri="other/doc3.txt")
# Count all
assert await client.count_documents() == 3
# Count with filter
assert await client.count_documents(filter="uri LIKE '%.pdf'") == 2
assert await client.count_documents(filter="uri LIKE '%.txt'") == 1
def test_citation_index_fallback_without_session_state():
"""Test that citation indices fall back to sequential numbering without session_state.
This tests the fallback branch in the ask and search tools when
ctx.deps.session_state is None.
"""
# Simulate the fallback logic from agent.py lines 281-285
citation_infos = []
session_state = None # No session state
# Simulate processing citations without session_state
chunk_ids = ["chunk-a", "chunk-b", "chunk-c"]
for chunk_id in chunk_ids:
if session_state is not None:
index = session_state.get_or_assign_index(chunk_id)
else:
index = len(citation_infos) + 1
citation_infos.append(
Citation(
index=index,
document_id="doc-1",
chunk_id=chunk_id,
document_uri="test.md",
content="test",
)
)
# Without session_state, indices are simple sequential numbers
assert citation_infos[0].index == 1
assert citation_infos[1].index == 2
assert citation_infos[2].index == 3
@pytest.mark.asyncio
async def test_summarization_task_cleanup_on_completion():
"""Test that completed tasks are cleaned up from _summarization_tasks."""
import asyncio
from haiku.rag.agents.chat.context import _summarization_tasks
_summarization_tasks.clear()
key = 67890 # Simulates id(qa_session_state)
# Create a fast task
async def fast_task():
await asyncio.sleep(0.01)
task = asyncio.create_task(fast_task())
_summarization_tasks[key] = task
task.add_done_callback(lambda t: _summarization_tasks.pop(key, None))
# Wait for completion
await task
# Task should be cleaned up
assert key not in _summarization_tasks
@pytest.mark.asyncio
async def test_ask_tool_returns_qa_result_when_no_state_delta(temp_db_path):
"""Test ask tool falls through to QAResult when state delta is empty.
When run_qa_core produces no state changes (client_snapshot == new_snapshot),
compute_combined_state_delta returns None and the tool returns QAResult directly.
"""
from unittest.mock import AsyncMock, patch
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from haiku.rag.tools.models import QAResult
from haiku.rag.tools.qa import create_qa_toolset
async with HaikuRAG(temp_db_path, create=True) as client:
context = ToolContext()
prepare_chat_context(context, features=["search", "qa"])
context.state_key = AGUI_STATE_KEY
deps = ChatDeps(config=Config, client=client, tool_context=context)
toolset = create_qa_toolset(Config)
agent = Agent(
TestModel(call_tools=["ask"]),
deps_type=ChatDeps,
toolsets=[toolset], # ty: ignore[invalid-argument-type]
)
# Mock run_qa_core to return a result WITHOUT modifying any state.
# Since client_snapshot == new_snapshot, the delta is None.
mock_result = QAResult(
question="test", answer="The answer", confidence=0.9, citations=[]
)
with patch(
"haiku.rag.tools.qa.run_qa_core",
new_callable=AsyncMock,
return_value=mock_result,
):
result = await agent.run(
"test question",
deps=deps, # ty: ignore[invalid-argument-type]
)
# The tool should have returned the QAResult (not a ToolReturn)
assert "The answer" in result.output
@pytest.mark.asyncio
async def test_ask_tool_delta_without_citations(temp_db_path):
"""Test ask tool emits StateDeltaEvent without citation sources when citations are empty.
When the QA result has no citations but state DID change,
the tool returns a ToolReturn with a delta but no 'Sources:' line.
"""
from unittest.mock import AsyncMock, patch
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from haiku.rag.tools.models import QAResult
from haiku.rag.tools.qa import (
QA_SESSION_NAMESPACE,
QASessionState,
create_qa_toolset,
)
async with HaikuRAG(temp_db_path, create=True) as client:
context = ToolContext()
prepare_chat_context(context, features=["search", "qa"])
context.state_key = AGUI_STATE_KEY
deps = ChatDeps(config=Config, client=client, tool_context=context)
toolset = create_qa_toolset(Config)
agent = Agent(
TestModel(call_tools=["ask"]),
deps_type=ChatDeps,
toolsets=[toolset], # ty: ignore[invalid-argument-type]
)
mock_result = QAResult(
question="test", answer="No citations answer", confidence=0.9, citations=[]
)
async def mock_qa_core(*args, **kwargs):
# Simulate state mutation (qa_history entry) but no citations
ctx = kwargs.get("context")
if ctx is not None:
qa_state = ctx.get(QA_SESSION_NAMESPACE, QASessionState)
if qa_state is not None:
qa_state.qa_history.append(
QAHistoryEntry(question="test", answer="No citations answer")
)
return mock_result
with patch(
"haiku.rag.tools.qa.run_qa_core",
new_callable=AsyncMock,
side_effect=mock_qa_core,
):
result = await agent.run(
"test question",
deps=deps, # ty: ignore[invalid-argument-type]
)
# Should have the answer but NO "Sources:" line
assert "No citations answer" in result.output
assert "Sources:" not in result.output
# Should have emitted a StateDeltaEvent (qa_history changed)
state = extract_state_from_result(result)
assert state is not None
assert len(state.get("qa_history", [])) > 0
@pytest.mark.asyncio
async def test_ask_tool_delta_with_citations(temp_db_path):
"""Test ask tool emits StateDeltaEvent with 'Sources:' line when citations are present.
When the QA result has citations and state changed, the tool returns
a ToolReturn with a delta and appended 'Sources: [1] [2]' line.
"""
from unittest.mock import AsyncMock, patch
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from haiku.rag.tools.models import QAResult
from haiku.rag.tools.qa import (
QA_SESSION_NAMESPACE,
QASessionState,
create_qa_toolset,
)
async with HaikuRAG(temp_db_path, create=True) as client:
context = ToolContext()
prepare_chat_context(context, features=["search", "qa"])
context.state_key = AGUI_STATE_KEY
deps = ChatDeps(config=Config, client=client, tool_context=context)
toolset = create_qa_toolset(Config)
agent = Agent(
TestModel(call_tools=["ask"]),
deps_type=ChatDeps,
toolsets=[toolset], # ty: ignore[invalid-argument-type]
)
citations = [
Citation(
index=1,
document_id="doc-1",
chunk_id="chunk-a",
document_uri="test.md",
document_title="Test",
content="content",
),
Citation(
index=2,
document_id="doc-1",
chunk_id="chunk-b",
document_uri="test.md",
document_title="Test",
content="content",
),
]
mock_result = QAResult(
question="test",
answer="Cited answer",
confidence=0.95,
citations=citations,
)
async def mock_qa_core(*args, **kwargs):
ctx = kwargs.get("context")
if ctx is not None:
qa_state = ctx.get(QA_SESSION_NAMESPACE, QASessionState)
if qa_state is not None:
qa_state.qa_history.append(
QAHistoryEntry(
question="test",
answer="Cited answer",
citations=citations,
)
)
return mock_result
with patch(
"haiku.rag.tools.qa.run_qa_core",
new_callable=AsyncMock,
side_effect=mock_qa_core,
):
result = await agent.run(
"test question",
deps=deps, # ty: ignore[invalid-argument-type]
)
# Should have the answer WITH "Sources:" line
assert "Cited answer" in result.output
assert "Sources:" in result.output
assert "[1]" in result.output
assert "[2]" in result.output
# Should have emitted a StateDeltaEvent
state = extract_state_from_result(result)
assert state is not None
Reference in a new issue View git blame Copy permalink