Improve coverage

This commit is contained in:
Yiorgis Gozadinos 2026-01-26 16:22:18 +02:00
parent 5976cebe8c
commit 6f9bc369dc
No known key found for this signature in database
4 changed files with 4550 additions and 1 deletions

View file

@ -123,7 +123,7 @@ state = ChatSessionState()
Q/A history is used to: Q/A history is used to:
1. Provide context for follow-up questions 1. Provide context for follow-up questions
2. Avoid repeating previous answers via the `recall` tool 2. Avoid repeating previous answers (the `ask` tool automatically recalls relevant prior answers)
3. Enable semantic ranking of relevant past answers 3. Enable semantic ranking of relevant past answers
### AG-UI Integration ### AG-UI Integration

View file

@ -515,6 +515,75 @@ async def test_chat_agent_ask_triggers_background_summarization(
assert session_state.session_context.last_updated is not None assert session_state.session_context.last_updated is not None
@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 (agent.py lines 231-257):
1. First ask populates qa_history with question_embedding
2. Second similar ask should find the prior answer via embedding similarity
"""
import asyncio
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",
)
agent = create_chat_agent(Config)
session_state = ChatSessionState(session_id="test-prior-answers")
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
)
# First ask - establishes qa_history
result1 = await agent.run(
"What are the class labels in DocLayNet?",
deps=deps,
)
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
# Wait for background summarization to complete
for _ in range(50):
if session_state.session_context is not None:
break
await asyncio.sleep(0.1)
# 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,
)
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
# 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"
)
def test_fifo_limit_enforcement(): def test_fifo_limit_enforcement():
"""Test that FIFO limit enforcement logic works correctly. """Test that FIFO limit enforcement logic works correctly.
@ -897,6 +966,39 @@ async def test_summarization_task_cancellation():
_summarization_tasks.clear() _summarization_tasks.clear()
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 @pytest.mark.asyncio
async def test_summarization_task_cleanup_on_completion(): async def test_summarization_task_cleanup_on_completion():
"""Test that completed tasks are cleaned up from _summarization_tasks.""" """Test that completed tasks are cleaned up from _summarization_tasks."""

View file

@ -7,6 +7,7 @@ from evaluations.evaluators import LLMJudge
from haiku.rag.agents.qa.agent import QuestionAnswerAgent from haiku.rag.agents.qa.agent import QuestionAnswerAgent
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.config.models import ModelConfig from haiku.rag.config.models import ModelConfig
HAS_ANTHROPIC = importlib.util.find_spec("anthropic") is not None HAS_ANTHROPIC = importlib.util.find_spec("anthropic") is not None
@ -17,6 +18,38 @@ def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_qa") return str(Path(__file__).parent.parent.parent / "cassettes" / "test_qa")
def test_get_qa_agent_factory(temp_db_path):
"""Test get_qa_agent factory function creates a properly configured agent."""
from haiku.rag.agents.qa import get_qa_agent
client = HaikuRAG(temp_db_path, create=True)
agent = get_qa_agent(client, Config)
assert agent is not None
assert isinstance(agent, QuestionAnswerAgent)
# Verify internal client is set correctly
assert agent._client is client
client.close()
def test_get_qa_agent_with_custom_prompt(temp_db_path):
"""Test get_qa_agent factory with custom system prompt."""
from haiku.rag.agents.qa import get_qa_agent
client = HaikuRAG(temp_db_path, create=True)
custom_prompt = "You are a custom QA assistant."
agent = get_qa_agent(client, Config, system_prompt=custom_prompt)
assert agent is not None
assert isinstance(agent, QuestionAnswerAgent)
# The internal pydantic-ai agent should have instructions set
# (pydantic-ai wraps the string in an Instructions object)
assert agent._agent.instructions is not None
client.close()
@pytest.mark.vcr() @pytest.mark.vcr()
async def test_qa_ollama(allow_model_requests, qa_corpus: Dataset, temp_db_path): async def test_qa_ollama(allow_model_requests, qa_corpus: Dataset, temp_db_path):
"""Test Ollama QA with LLM judge (VCR recorded).""" """Test Ollama QA with LLM judge (VCR recorded)."""

File diff suppressed because one or more lines are too long