Improve coverage
This commit is contained in:
parent
5976cebe8c
commit
6f9bc369dc
4 changed files with 4550 additions and 1 deletions
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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."""
|
||||||
|
|
|
||||||
|
|
@ -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
Loading…
Reference in a new issue