480 lines
15 KiB
Python
480 lines
15 KiB
Python
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from haiku.rag.agents.chat import (
|
|
ChatDeps,
|
|
ChatSessionState,
|
|
CitationInfo,
|
|
QAResponse,
|
|
SearchAgent,
|
|
create_chat_agent,
|
|
)
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import Config
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def vcr_cassette_dir():
|
|
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_chat_agent")
|
|
|
|
|
|
def test_create_chat_agent():
|
|
"""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)
|
|
deps = ChatDeps(client=client, config=Config)
|
|
|
|
assert deps.client is client
|
|
assert deps.config is Config
|
|
assert deps.search_results is None
|
|
assert deps.session_state is None
|
|
|
|
client.close()
|
|
|
|
|
|
def test_chat_session_state():
|
|
"""Test ChatSessionState model."""
|
|
state = ChatSessionState(session_id="test-session")
|
|
assert state.session_id == "test-session"
|
|
assert state.citations == []
|
|
assert state.qa_history == []
|
|
|
|
|
|
def test_citation_info():
|
|
"""Test CitationInfo model."""
|
|
citation = CitationInfo(
|
|
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 QAResponse model."""
|
|
citation = CitationInfo(
|
|
index=1,
|
|
document_id="doc-123",
|
|
chunk_id="chunk-456",
|
|
document_uri="test.md",
|
|
document_title="Test Document",
|
|
content="Test content",
|
|
)
|
|
qa = QAResponse(
|
|
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 QAResponse.sources falls back to URI when title is None."""
|
|
citation = CitationInfo(
|
|
index=1,
|
|
document_id="doc-123",
|
|
chunk_id="chunk-456",
|
|
document_uri="test.md",
|
|
document_title=None,
|
|
content="Test content",
|
|
)
|
|
qa = QAResponse(
|
|
question="What is this?",
|
|
answer="This is a test",
|
|
citations=[citation],
|
|
)
|
|
assert qa.sources == ["test.md"]
|
|
|
|
|
|
def test_search_agent_initialization(temp_db_path):
|
|
"""Test SearchAgent can be initialized."""
|
|
client = HaikuRAG(temp_db_path, create=True)
|
|
search_agent = SearchAgent(client, Config)
|
|
assert search_agent is not None
|
|
client.close()
|
|
|
|
|
|
# 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_chat_agent_with_qa_history_ranking(allow_model_requests, temp_db_path):
|
|
"""Test chat agent uses similarity ranking for qa_history.
|
|
|
|
This test verifies that when qa_history has more than 5 entries,
|
|
the ranking function is applied and the agent can still process requests.
|
|
"""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
# Add a simple document
|
|
await client.create_document(
|
|
content=DOCLAYNET_CLASS_LABELS,
|
|
uri="doclaynet-labels",
|
|
title="DocLayNet Class Labels",
|
|
)
|
|
|
|
agent = create_chat_agent(Config)
|
|
|
|
# Build session state with pre-populated qa_history (>5 items to trigger ranking)
|
|
session_state = ChatSessionState(
|
|
session_id="test-ranking",
|
|
qa_history=[
|
|
QAResponse(
|
|
question="What are the 11 class labels in DocLayNet?",
|
|
answer="The 11 class labels are: Caption, Footnote, Formula, List-item, Page-footer, Page-header, Picture, Section-header, Table, Text, and Title.",
|
|
),
|
|
QAResponse(
|
|
question="How was the annotation process organized?",
|
|
answer="The annotation was organized into 4 phases.",
|
|
),
|
|
QAResponse(
|
|
question="What data sources were used?",
|
|
answer="Sources include arXiv and government offices.",
|
|
),
|
|
QAResponse(
|
|
question="How were pages selected?",
|
|
answer="By selective subsampling.",
|
|
),
|
|
QAResponse(
|
|
question="What is the agreement metric?",
|
|
answer="The mAP metric was used.",
|
|
),
|
|
QAResponse(
|
|
question="What is machine learning?",
|
|
answer="A field of AI.",
|
|
),
|
|
],
|
|
)
|
|
|
|
deps = ChatDeps(
|
|
client=client,
|
|
config=Config,
|
|
session_state=session_state,
|
|
)
|
|
|
|
# Ask a question - the key test is that ranking is applied without error
|
|
result = await agent.run(
|
|
"What class labels are defined in the dataset?",
|
|
deps=deps,
|
|
)
|
|
|
|
# Verify the agent produced a response (ranking didn't break anything)
|
|
assert result.output is not None
|
|
assert len(result.output) > 0
|
|
|
|
# Verify qa_history was updated (new Q&A was added)
|
|
assert len(session_state.qa_history) == 7 # 6 original + 1 new
|
|
|
|
|
|
@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",
|
|
)
|
|
|
|
agent = create_chat_agent(Config)
|
|
session_state = ChatSessionState(session_id="test-search")
|
|
deps = ChatDeps(
|
|
client=client,
|
|
config=Config,
|
|
session_state=session_state,
|
|
)
|
|
|
|
# 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",
|
|
)
|
|
|
|
agent = create_chat_agent(Config)
|
|
session_state = ChatSessionState(session_id="test-search-filter")
|
|
deps = ChatDeps(
|
|
client=client,
|
|
config=Config,
|
|
session_state=session_state,
|
|
)
|
|
|
|
# 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",
|
|
)
|
|
|
|
agent = create_chat_agent(Config)
|
|
deps = ChatDeps(
|
|
client=client,
|
|
config=Config,
|
|
)
|
|
|
|
# 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:
|
|
agent = create_chat_agent(Config)
|
|
deps = ChatDeps(
|
|
client=client,
|
|
config=Config,
|
|
)
|
|
|
|
# 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_search_agent_with_context(allow_model_requests, temp_db_path):
|
|
"""Test SearchAgent's search method with context."""
|
|
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",
|
|
)
|
|
|
|
search_agent = SearchAgent(client, Config)
|
|
|
|
# Search with context
|
|
results = await search_agent.search(
|
|
query="What are the class labels?",
|
|
context="We're discussing document layout analysis",
|
|
)
|
|
|
|
assert isinstance(results, list)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_search_agent_with_filter(allow_model_requests, temp_db_path):
|
|
"""Test SearchAgent's search method 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",
|
|
)
|
|
|
|
search_agent = SearchAgent(client, Config)
|
|
|
|
# Search with filter - only the labels document
|
|
results = await search_agent.search(
|
|
query="What information is available?",
|
|
filter="uri LIKE '%labels%'",
|
|
)
|
|
|
|
assert isinstance(results, list)
|
|
# Results should only come from the labels document
|
|
for r in results:
|
|
assert "labels" in (r.document_uri or "")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_search_agent_deduplication(allow_model_requests, temp_db_path):
|
|
"""Test SearchAgent deduplicates results by chunk_id."""
|
|
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",
|
|
)
|
|
|
|
search_agent = SearchAgent(client, Config)
|
|
|
|
# Search - the search agent will likely run multiple queries
|
|
# that could return the same chunk, which should be deduplicated
|
|
results = await search_agent.search(
|
|
query="Tell me about class labels and their counts",
|
|
)
|
|
|
|
assert isinstance(results, list)
|
|
|
|
# Verify no duplicate chunk_ids
|
|
chunk_ids = [r.chunk_id for r in results if r.chunk_id]
|
|
assert len(chunk_ids) == len(set(chunk_ids)), "Found duplicate chunk_ids"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_search_agent_no_results(allow_model_requests, temp_db_path):
|
|
"""Test SearchAgent handles no results gracefully."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
search_agent = SearchAgent(client, Config)
|
|
|
|
# Search in empty database
|
|
results = await search_agent.search(
|
|
query="Find information about nonexistent topic xyz123",
|
|
)
|
|
|
|
assert isinstance(results, list)
|
|
assert len(results) == 0
|
|
|
|
|
|
@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 adds citations to the response."""
|
|
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",
|
|
)
|
|
|
|
agent = create_chat_agent(Config)
|
|
session_state = ChatSessionState(session_id="test-citations")
|
|
deps = ChatDeps(
|
|
client=client,
|
|
config=Config,
|
|
session_state=session_state,
|
|
)
|
|
|
|
# Ask a question that should use the ask tool with citations
|
|
result = await agent.run(
|
|
"What is the highest count class in the DocLayNet dataset?",
|
|
deps=deps,
|
|
)
|
|
|
|
assert result.output is not None
|
|
# The qa_history should have been updated with the new Q&A
|
|
assert len(session_state.qa_history) >= 1
|