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() )