From 70544b5112cbe0514be081d3b38b06822f49f57a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 16 Sep 2025 15:13:52 +0300 Subject: [PATCH] Research agent tests --- tests/research/test_evaluation_agent.py | 14 ++ tests/research/test_orchestrator.py | 102 ++++++-------- tests/research/test_search_agent.py | 174 +----------------------- tests/research/test_synthesis_agent.py | 11 ++ 4 files changed, 72 insertions(+), 229 deletions(-) create mode 100644 tests/research/test_evaluation_agent.py create mode 100644 tests/research/test_synthesis_agent.py diff --git a/tests/research/test_evaluation_agent.py b/tests/research/test_evaluation_agent.py new file mode 100644 index 00000000..fd63c5a4 --- /dev/null +++ b/tests/research/test_evaluation_agent.py @@ -0,0 +1,14 @@ +from haiku.rag.research.evaluation_agent import ( + AnalysisEvaluationAgent, + EvaluationResult, +) + + +class TestAnalysisEvaluationAgent: + """Lean tests for AnalysisEvaluationAgent without LLM mocking.""" + + def test_agent_initialization(self): + agent = AnalysisEvaluationAgent(provider="openai", model="gpt-4") + assert agent.provider == "openai" + assert agent.model == "gpt-4" + assert agent.output_type == EvaluationResult diff --git a/tests/research/test_orchestrator.py b/tests/research/test_orchestrator.py index 635bbaea..09d44c6e 100644 --- a/tests/research/test_orchestrator.py +++ b/tests/research/test_orchestrator.py @@ -1,16 +1,22 @@ -"""Tests for the research orchestrator.""" - -from unittest.mock import AsyncMock, MagicMock, create_autospec +from unittest.mock import AsyncMock, create_autospec import pytest +from pydantic_ai.models.test import TestModel from haiku.rag.client import HaikuRAG from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies from haiku.rag.research.evaluation_agent import EvaluationResult from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan +from haiku.rag.research.synthesis_agent import ResearchReport from haiku.rag.store.models.chunk import Chunk +@pytest.fixture +def test_model(): + """Create a test model for orchestrator testing.""" + return TestModel() + + @pytest.fixture def mock_client(): """Create a mock HaikuRAG client.""" @@ -100,6 +106,9 @@ class TestResearchOrchestrator: orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4") # Create mock evaluation results + from unittest.mock import MagicMock + + # Sufficient research result sufficient_result = MagicMock() sufficient_result.output = EvaluationResult( key_insights=["Climate is changing", "Human activity is the cause"], @@ -109,6 +118,7 @@ class TestResearchOrchestrator: reasoning="All aspects covered comprehensively", ) + # Insufficient research result insufficient_result = MagicMock() insufficient_result.output = EvaluationResult( key_insights=["Some data found"], @@ -136,66 +146,34 @@ class TestResearchOrchestrator: assert not orchestrator._should_stop_research(insufficient_result, 0.8) @pytest.mark.asyncio - async def test_conduct_research_workflow(self, mock_client): - """Test the basic research workflow.""" + async def test_conduct_research_workflow(self, test_model, mock_client): + """Test the basic research workflow using TestModel.""" orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4") - # Mock the agent runs - # Mock initial plan - plan_mock = MagicMock() - plan_mock.output = ResearchPlan( - main_question="What is climate change?", - sub_questions=[ - "What causes climate change?", - "What are the effects?", - "What can be done?", - ], - ) - orchestrator.run = AsyncMock(return_value=plan_mock) + # Setup mock client returns + mock_chunks = [ + create_mock_chunk("1", "Climate change information"), + ] + mock_client.search.return_value = mock_chunks + mock_client.expand_context.return_value = mock_chunks - # Mock search agent - search_mock = MagicMock() - search_mock.output = "Climate change is caused by greenhouse gases." - orchestrator.search_agent.run = AsyncMock(return_value=search_mock) + # Use TestModel for all agents + with orchestrator.agent.override(model=test_model): + with orchestrator.search_agent.agent.override(model=test_model): + with orchestrator.evaluation_agent.agent.override(model=test_model): + with orchestrator.synthesis_agent.agent.override(model=test_model): + # Run the research + report = await orchestrator.conduct_research( + "What is climate change?", mock_client, max_iterations=1 + ) - # Mock evaluation agent - make it stop after first iteration - eval_mock = MagicMock() - eval_mock.output = EvaluationResult( - key_insights=["Climate change is real"], - new_questions=[], - confidence_score=0.9, - is_sufficient=True, - reasoning="Sufficient information gathered", - ) - orchestrator.evaluation_agent.run = AsyncMock(return_value=eval_mock) - - # Mock synthesis agent - from haiku.rag.research.synthesis_agent import ResearchReport - - synthesis_mock = MagicMock() - synthesis_mock.output = ResearchReport( - title="Climate Change Report", - executive_summary="Summary", - main_findings=["Finding 1"], - themes={}, - conclusions=[], - limitations=[], - recommendations=[], - sources_summary="Sources", - ) - orchestrator.synthesis_agent.run = AsyncMock(return_value=synthesis_mock) - - # Mock client search and expand - mock_client.search.return_value = [] - mock_client.expand_context.return_value = [] - - # Run the research - report = await orchestrator.conduct_research( - "What is climate change?", mock_client, max_iterations=3 - ) - - # Verify we got a report - assert report.title == "Climate Change Report" - - # Verify search was called for all 3 sub-questions - assert orchestrator.search_agent.run.call_count == 3 + # Verify we got a valid report structure + assert isinstance(report, ResearchReport) + assert report.title + assert report.executive_summary + assert isinstance(report.main_findings, list) + assert isinstance(report.themes, dict) + assert isinstance(report.conclusions, list) + assert isinstance(report.limitations, list) + assert isinstance(report.recommendations, list) + assert report.sources_summary diff --git a/tests/research/test_search_agent.py b/tests/research/test_search_agent.py index 6105c9f6..d6d18b1e 100644 --- a/tests/research/test_search_agent.py +++ b/tests/research/test_search_agent.py @@ -1,171 +1,11 @@ -"""Tests for the search specialist agent.""" - -from unittest.mock import AsyncMock, create_autospec - -import pytest -from pydantic_ai import RunContext -from pydantic_ai.models.test import TestModel -from pydantic_ai.usage import RunUsage - -from haiku.rag.client import HaikuRAG -from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies from haiku.rag.research.search_agent import SearchSpecialistAgent -from haiku.rag.store.models.chunk import Chunk -@pytest.fixture -def mock_client(): - """Create a mock HaikuRAG client.""" - client = create_autospec(HaikuRAG, instance=True) - client.search = AsyncMock() - client.expand_context = AsyncMock() - return client +class TestSearchSpecialistAgent: + """Lean tests for SearchSpecialistAgent without LLM mocking.""" - -@pytest.fixture -def research_context(): - """Create a research context for testing.""" - return ResearchContext( - original_question="What is climate change?", - ) - - -@pytest.fixture -def research_deps(mock_client, research_context): - """Create research dependencies for testing.""" - return ResearchDependencies(client=mock_client, context=research_context) - - -def create_mock_chunk(chunk_id: str, content: str, score: float = 0.8): - """Helper to create mock chunk objects.""" - return Chunk( - id=chunk_id, - document_id=f"doc_{chunk_id}", - content=content, - document_uri=f"doc_{chunk_id}.md", - metadata={}, - ), score - - -def get_agent_tool(agent, tool_name: str): - """Helper to get a tool from an agent by name.""" - tools = agent.agent._function_toolset.tools - if tool_name in tools: - return tools[tool_name].function - return None - - -@pytest.mark.asyncio -async def test_search_agent_has_search_tool(): - """Test that the search agent registers a search tool.""" - test_model = TestModel() - # Use a valid provider for initialization - agent = SearchSpecialistAgent(provider="openai", model="gpt-4") - - # Run agent with TestModel to check tools - with agent.agent.override(model=test_model): - await agent.agent.run( - "test", - deps=ResearchDependencies( - client=create_autospec(HaikuRAG, instance=True), - context=ResearchContext(original_question="test"), - ), - ) - - # Verify the search tool was registered - assert test_model.last_model_request_parameters is not None - tools = test_model.last_model_request_parameters.function_tools - assert tools is not None - assert len(tools) == 1 - assert tools[0].name == "search_and_answer" - - -@pytest.mark.asyncio -async def test_search_single_query(mock_client, research_deps): - """Test that search tool is called with single query.""" - # Setup mock responses - mock_chunks = [ - create_mock_chunk("chunk1", "Climate change is a global phenomenon"), - create_mock_chunk("chunk2", "Rising temperatures affect ecosystems"), - ] - - mock_client.search.return_value = mock_chunks[:1] - mock_client.expand_context.return_value = mock_chunks - - # Create agent - agent = SearchSpecialistAgent(provider="openai", model="gpt-4") - - # Get the search tool - search_tool = get_agent_tool(agent, "search_and_answer") - assert search_tool is not None - - # Test the tool - ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage()) - result = await search_tool(ctx, query="climate change") - - # Verify result - should be a formatted string with context - assert isinstance(result, str) - assert "Climate change is a global phenomenon" in result - assert "Rising temperatures affect ecosystems" in result - - # Verify mock was called with default limit - mock_client.search.assert_called_once_with("climate change", limit=5) - mock_client.expand_context.assert_called_once() - - -@pytest.mark.asyncio -async def test_search_with_limit(mock_client, research_deps): - """Test that search respects the limit parameter.""" - # Create more chunks than the limit - mock_chunks = [ - create_mock_chunk("chunk1", "Content 1", 0.9), - create_mock_chunk("chunk2", "Content 2", 0.8), - create_mock_chunk("chunk3", "Content 3", 0.7), - create_mock_chunk("chunk4", "Content 4", 0.6), - create_mock_chunk("chunk5", "Content 5", 0.5), - ] - - mock_client.search.return_value = mock_chunks[:3] - mock_client.expand_context.return_value = mock_chunks[:3] - - agent = SearchSpecialistAgent(provider="openai", model="gpt-4") - - # Get the search tool - search_tool = get_agent_tool(agent, "search_and_answer") - assert search_tool is not None - - # Test the tool with limit - ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage()) - result = await search_tool(ctx, query="test query", limit=3) - - # Verify result is a formatted string - assert isinstance(result, str) - assert "Content 1" in result - assert "Content 2" in result - assert "Content 3" in result - - # Verify mock was called with correct limit - mock_client.search.assert_called_once_with("test query", limit=3) - - -@pytest.mark.asyncio -async def test_search_updates_context(mock_client, research_deps): - """Test that search results are stored in context.""" - mock_chunks = [create_mock_chunk("chunk1", "Test content")] - - mock_client.search.return_value = [] - mock_client.expand_context.return_value = mock_chunks - - agent = SearchSpecialistAgent(provider="openai", model="gpt-4") - - # Get the search tool - search_tool = get_agent_tool(agent, "search_and_answer") - assert search_tool is not None - - # Test the tool - ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage()) - await search_tool(ctx, query="test query") - - # Verify context was updated - assert len(research_deps.context.search_results) == 1 - assert research_deps.context.search_results[0]["query"] == "test query" + def test_agent_initialization(self): + agent = SearchSpecialistAgent(provider="openai", model="gpt-4") + assert agent.provider == "openai" + assert agent.model == "gpt-4" + assert agent.output_type is str diff --git a/tests/research/test_synthesis_agent.py b/tests/research/test_synthesis_agent.py new file mode 100644 index 00000000..666662c7 --- /dev/null +++ b/tests/research/test_synthesis_agent.py @@ -0,0 +1,11 @@ +from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent + + +class TestSynthesisAgent: + """Lean tests for SynthesisAgent without LLM mocking.""" + + def test_agent_initialization(self): + agent = SynthesisAgent(provider="openai", model="gpt-4") + assert agent.provider == "openai" + assert agent.model == "gpt-4" + assert agent.output_type == ResearchReport