Research agent tests

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

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

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

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