258 lines
9.1 KiB
Python
258 lines
9.1 KiB
Python
from pathlib import Path
|
|
|
|
import pytest
|
|
from pydantic_ai import Agent
|
|
|
|
from haiku.rag.agents.analysis.agent import create_analysis_agent
|
|
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
|
|
from haiku.rag.agents.analysis.models import CodeExecution, RawAnalysisResult
|
|
from haiku.rag.config import AppConfig, Config
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def vcr_cassette_dir():
|
|
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_analysis")
|
|
|
|
|
|
class TestCreateAnalysisAgent:
|
|
def test_creates_agent(self):
|
|
agent = create_analysis_agent(Config)
|
|
assert isinstance(agent, Agent)
|
|
assert agent.deps_type is AnalysisDeps
|
|
assert agent.output_type is RawAnalysisResult
|
|
|
|
def test_agent_has_execute_code_tool(self):
|
|
agent = create_analysis_agent(Config)
|
|
tool_names = list(agent._function_toolset.tools.keys())
|
|
assert "execute_code" in tool_names
|
|
|
|
|
|
class TestCodeExecutionModel:
|
|
def test_code_execution_has_correct_fields(self):
|
|
"""Test that CodeExecution has all expected fields."""
|
|
execution = CodeExecution(
|
|
code="print('hello')",
|
|
stdout="hello\n",
|
|
stderr="",
|
|
success=True,
|
|
)
|
|
assert execution.code == "print('hello')"
|
|
assert execution.stdout == "hello\n"
|
|
assert execution.stderr == ""
|
|
assert execution.success is True
|
|
|
|
|
|
class TestClientAnalysisIntegration:
|
|
"""Integration tests for client.analyze() method."""
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_analyze_count_documents(self, allow_model_requests, temp_db_path):
|
|
"""Test analysis agent can count documents.
|
|
|
|
Agent program:
|
|
docs = list_documents(limit=1000)
|
|
print(len(docs))
|
|
"""
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
config = AppConfig()
|
|
|
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
|
await client.create_document("First document about cats.", title="Doc 1")
|
|
await client.create_document("Second document about dogs.", title="Doc 2")
|
|
await client.create_document("Third document about birds.", title="Doc 3")
|
|
|
|
result = await client.analyze("How many documents are in the database?")
|
|
|
|
assert "3" in result.answer
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_analyze_aggregation(self, allow_model_requests, temp_db_path):
|
|
"""Test analysis agent can perform aggregation across documents.
|
|
|
|
Agent program:
|
|
import re
|
|
revs = {}
|
|
for d in ['Q1 Report', 'Q2 Report', 'Q3 Report']:
|
|
content = get_document(d)
|
|
if content:
|
|
vals = re.findall(r'\\$([\\d,]+)', content)
|
|
if vals:
|
|
rev = sum(int(v.replace(',', '')) for v in vals)
|
|
else:
|
|
rev = None
|
|
else:
|
|
rev = None
|
|
revs[d] = rev
|
|
print(revs)
|
|
"""
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
config = AppConfig()
|
|
|
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
|
await client.create_document(
|
|
"Sales report Q1: Revenue was $100,000.", title="Q1 Report"
|
|
)
|
|
await client.create_document(
|
|
"Sales report Q2: Revenue was $150,000.", title="Q2 Report"
|
|
)
|
|
await client.create_document(
|
|
"Sales report Q3: Revenue was $200,000.", title="Q3 Report"
|
|
)
|
|
|
|
result = await client.analyze(
|
|
"What is the total revenue across all quarterly reports?"
|
|
)
|
|
|
|
assert "450" in result.answer or "450,000" in result.answer
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_analyze_with_filter(self, allow_model_requests, temp_db_path):
|
|
"""Test analysis agent respects filter parameter.
|
|
|
|
Agent program:
|
|
docs = list_documents(limit=1000)
|
|
print(len(docs))
|
|
print(docs[:5])
|
|
|
|
The filter is applied via context, so list_documents() only sees "Cats".
|
|
"""
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
config = AppConfig()
|
|
|
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
|
await client.create_document("Cat document.", title="Cats")
|
|
await client.create_document("Dog document.", title="Dogs")
|
|
await client.create_document("Bird document.", title="Birds")
|
|
|
|
result = await client.analyze(
|
|
"How many documents are available?",
|
|
filter="title = 'Cats'",
|
|
)
|
|
|
|
assert "1" in result.answer
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_analyze_search_and_identify_source(
|
|
self, allow_model_requests, temp_db_path
|
|
):
|
|
"""Test analysis agent can search and identify source documents."""
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
config = AppConfig()
|
|
|
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
|
await client.create_document(
|
|
"The quick brown fox jumps over the lazy dog.",
|
|
title="Animal Facts",
|
|
)
|
|
|
|
result = await client.analyze(
|
|
"Search for content about animals and tell me "
|
|
"which document it came from."
|
|
)
|
|
|
|
assert "Animal Facts" in result.answer
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_analyze_search_and_extract(self, allow_model_requests, temp_db_path):
|
|
"""Test analysis agent can use search() to find content and extract information.
|
|
|
|
Agent program:
|
|
results = search("document element types", limit=20)
|
|
print(len(results))
|
|
for r in results[:5]:
|
|
print(r['document_title'], r['chunk_id'], r['score'])
|
|
print(r['content'][:200])
|
|
|
|
results = search("DocBank element types", limit=10)
|
|
...
|
|
"""
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
pdf_path = Path("tests/data/doclaynet.pdf")
|
|
config = AppConfig()
|
|
config.processing.conversion_options.do_ocr = False
|
|
|
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
|
await client.create_document_from_source(pdf_path)
|
|
|
|
result = await client.analyze(
|
|
"Search for content about document element types or labels. "
|
|
"What are all the different document element types mentioned? "
|
|
"List them all."
|
|
)
|
|
|
|
# The doclaynet.pdf defines exactly 11 class labels for document elements
|
|
# Normalize Unicode hyphens (U+2011 non-breaking hyphen) to regular hyphens
|
|
answer_lower = result.answer.lower().replace("\u2011", "-")
|
|
expected_labels = [
|
|
"caption",
|
|
"footnote",
|
|
"formula",
|
|
"list-item",
|
|
"page-footer",
|
|
"page-header",
|
|
"picture",
|
|
"section-header",
|
|
"table",
|
|
"text",
|
|
"title",
|
|
]
|
|
# Check that the agent found at least 6 of the 11 labels
|
|
# (LLM summaries may not always include all labels)
|
|
found_labels = [
|
|
label
|
|
for label in expected_labels
|
|
if label in answer_lower or label.replace("-", " ") in answer_lower
|
|
]
|
|
assert len(found_labels) >= 6, (
|
|
f"Expected at least 6 labels, found {len(found_labels)}: {found_labels}"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_analyze_with_preloaded_documents(
|
|
self, allow_model_requests, temp_db_path
|
|
):
|
|
"""Test analysis agent can use pre-loaded documents variable.
|
|
|
|
Agent program:
|
|
if 'documents' in dir():
|
|
for doc in documents:
|
|
print(doc['title'], len(doc['content']))
|
|
else:
|
|
print('No preloaded documents')
|
|
"""
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
config = AppConfig()
|
|
|
|
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
|
await client.create_document(
|
|
"The company was founded in 1985 by Jane Smith.",
|
|
title="Company History",
|
|
)
|
|
await client.create_document(
|
|
"Our mission is to make technology accessible to everyone.",
|
|
title="Mission Statement",
|
|
)
|
|
|
|
result = await client.analyze(
|
|
"Using the pre-loaded documents variable, "
|
|
"tell me when was the company founded and what is their mission?",
|
|
documents=["Company History", "Mission Statement"],
|
|
)
|
|
|
|
assert "1985" in result.answer
|
|
assert (
|
|
"accessible" in result.answer.lower()
|
|
or "technology" in result.answer.lower()
|
|
)
|