Use TestModel from pydantic AI where appropriate

tests/test_client.py

@@ -538,33 +538,51 @@ async def test_client_create_document_with_custom_chunks(temp_db_path):
 
 
 @pytest.mark.asyncio
-async def test_client_ask_without_cite(temp_db_path):
+async def test_client_ask_without_cite(monkeypatch, temp_db_path):
     """Test asking questions without citations."""
+    from pydantic_ai.models.test import TestModel
+
+    # Mock OpenAIChatModel to return TestModel
+    monkeypatch.setattr(
+        "haiku.rag.qa.agent.OpenAIChatModel", lambda **kwargs: TestModel()
+    )
+
     async with HaikuRAG(temp_db_path) as client:
-        # Mock the QA agent
+        # Create a test document for the agent to search
-        mock_qa_agent = AsyncMock()
+        await client.create_document(
-        mock_qa_agent.answer.return_value = "Test answer"
+            content="Python is a high-level programming language.", uri="test.txt"
+        )
 
-        with patch("haiku.rag.qa.get_qa_agent", return_value=mock_qa_agent):
+        # Use real QA agent with TestModel
-            answer = await client.ask("What is Python?")
+        answer = await client.ask("What is Python?")
 
-        assert answer == "Test answer"
+        # TestModel will generate a valid string response
-        mock_qa_agent.answer.assert_called_once_with("What is Python?")
+        assert answer is not None
+        assert isinstance(answer, str)
 
 
 @pytest.mark.asyncio
-async def test_client_ask_with_cite(temp_db_path):
+async def test_client_ask_with_cite(monkeypatch, temp_db_path):
     """Test asking questions with citations."""
+    from pydantic_ai.models.test import TestModel
+
+    # Mock OpenAIChatModel to return TestModel
+    monkeypatch.setattr(
+        "haiku.rag.qa.agent.OpenAIChatModel", lambda **kwargs: TestModel()
+    )
+
     async with HaikuRAG(temp_db_path) as client:
-        # Mock the QA agent
+        # Create a test document
-        mock_qa_agent = AsyncMock()
+        await client.create_document(
-        mock_qa_agent.answer.return_value = "Test answer with citations [1]"
+            content="Python is a high-level programming language.", uri="test.txt"
+        )
 
-        with patch("haiku.rag.qa.get_qa_agent", return_value=mock_qa_agent):
+        # Use real QA agent with TestModel
-            answer = await client.ask("What is Python?", cite=True)
+        answer = await client.ask("What is Python?", cite=True)
 
-        assert answer == "Test answer with citations [1]"
+        # TestModel will generate a valid string response
-        mock_qa_agent.answer.assert_called_once_with("What is Python?")
+        assert answer is not None
+        assert isinstance(answer, str)
 
 
 @pytest.mark.asyncio

tests/test_deep_qa.py

@@ -1,22 +1,17 @@
-from typing import Any, cast
-
 import pytest
+from pydantic_ai.models.test import TestModel
 
+from haiku.rag.client import HaikuRAG
 from haiku.rag.graph.models import SearchAnswer
 from haiku.rag.qa.deep.dependencies import DeepQAContext
 from haiku.rag.qa.deep.graph import build_deep_qa_graph
-from haiku.rag.qa.deep.models import DeepQAAnswer
+from haiku.rag.qa.deep.nodes import DeepQAPlanNode
-from haiku.rag.qa.deep.nodes import (
-    DeepQADecisionNode,
-    DeepQAPlanNode,
-    DeepQASearchDispatchNode,
-    DeepQASynthesizeNode,
-)
 from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
 
 
 @pytest.mark.asyncio
-async def test_deep_qa_graph_end_to_end(monkeypatch):
+async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
+    """Test deep Q&A graph with mocked LLM using TestModel."""
     graph = build_deep_qa_graph()
 
     state = DeepQAState(
@@ -25,116 +20,59 @@ async def test_deep_qa_graph_end_to_end(monkeypatch):
         ),
         max_sub_questions=3,
     )
-    deps = DeepQADeps(client=cast(Any, None), console=None)
 
-    async def fake_plan_run(self, ctx) -> Any:
+    # Use real client but with TestModel for LLM calls
-        ctx.state.context.sub_questions = [
+    client = HaikuRAG(temp_db_path)
-            "Describe haiku.rag in one sentence",
+    deps = DeepQADeps(client=client, console=None)
-            "List core components of haiku.rag",
-        ]
-        return DeepQASearchDispatchNode(self.provider, self.model)
 
-    async def fake_search_dispatch_run(self, ctx) -> Any:
+    # Mock get_model to return TestModel which generates valid schema-compliant data
-        if not ctx.state.context.sub_questions:
+    def test_model_factory(provider, model):
-            return DeepQADecisionNode(self.provider, self.model)
+        return TestModel()
 
-        batch = ctx.state.context.sub_questions[:]
+    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
-        ctx.state.context.sub_questions.clear()
+    monkeypatch.setattr("haiku.rag.qa.deep.nodes.get_model", test_model_factory)
 
-        for question in batch:
+    start = DeepQAPlanNode(provider="test", model="test")
-            ctx.state.context.add_qa_response(
-                SearchAnswer(
-                    query=question,
-                    answer=f"Answer to: {question}",
-                    context=["Context snippet"],
-                    sources=["test.md"],
-                )
-            )
-        return DeepQASearchDispatchNode(self.provider, self.model)
-
-    async def fake_decision_run(self, ctx) -> Any:
-        ctx.state.iterations += 1
-        return DeepQASynthesizeNode(self.provider, self.model)
-
-    async def fake_synthesize_run(self, ctx) -> Any:
-        from pydantic_graph import End
-
-        return End(
-            DeepQAAnswer(
-                answer="haiku.rag is a RAG system with components A, B, C.",
-                sources=["test.md"],
-            )
-        )
-
-    monkeypatch.setattr(DeepQAPlanNode, "run", fake_plan_run)
-    monkeypatch.setattr(DeepQASearchDispatchNode, "run", fake_search_dispatch_run)
-    monkeypatch.setattr(DeepQADecisionNode, "run", fake_decision_run)
-    monkeypatch.setattr(DeepQASynthesizeNode, "run", fake_synthesize_run)
-
-    start = DeepQAPlanNode(provider="ollama", model="test")
     result = await graph.run(start_node=start, state=state, deps=deps)
 
-    assert result.output.answer == "haiku.rag is a RAG system with components A, B, C."
+    # TestModel will generate valid structured output based on schemas
-    assert result.output.sources == ["test.md"]
+    assert result.output.answer is not None
-    assert len(state.context.qa_responses) == 2
+    assert isinstance(result.output.answer, str)
+    assert isinstance(result.output.sources, list)
+
+    client.close()
 
 
 @pytest.mark.asyncio
-async def test_deep_qa_with_citations(monkeypatch):
+async def test_deep_qa_with_citations(monkeypatch, temp_db_path):
+    """Test deep Q&A with citations enabled using TestModel."""
     graph = build_deep_qa_graph()
 
     state = DeepQAState(
         context=DeepQAContext(original_question="What is Python?", use_citations=True),
         max_sub_questions=2,
     )
-    deps = DeepQADeps(client=cast(Any, None), console=None)
 
-    async def fake_plan_run(self, ctx) -> Any:
+    # Use real client but with TestModel for LLM calls
-        ctx.state.context.sub_questions = ["What is Python used for?"]
+    client = HaikuRAG(temp_db_path)
-        return DeepQASearchDispatchNode(self.provider, self.model)
+    deps = DeepQADeps(client=client, console=None)
 
-    async def fake_search_dispatch_run(self, ctx) -> Any:
+    # Mock get_model to return TestModel
-        if not ctx.state.context.sub_questions:
+    def test_model_factory(provider, model):
-            return DeepQADecisionNode(self.provider, self.model)
+        return TestModel()
 
-        batch = ctx.state.context.sub_questions[:]
+    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
-        ctx.state.context.sub_questions.clear()
+    monkeypatch.setattr("haiku.rag.qa.deep.nodes.get_model", test_model_factory)
 
-        for question in batch:
+    start = DeepQAPlanNode(provider="test", model="test")
-            ctx.state.context.add_qa_response(
-                SearchAnswer(
-                    query=question,
-                    answer="Python is used for web development and data science.",
-                    context=["Python snippet"],
-                    sources=["python.md"],
-                )
-            )
-        return DeepQASearchDispatchNode(self.provider, self.model)
-
-    async def fake_decision_run(self, ctx) -> Any:
-        ctx.state.iterations += 1
-        return DeepQASynthesizeNode(self.provider, self.model)
-
-    async def fake_synthesize_run(self, ctx) -> Any:
-        from pydantic_graph import End
-
-        return End(
-            DeepQAAnswer(
-                answer="Python is a programming language [python.md].",
-                sources=["python.md"],
-            )
-        )
-
-    monkeypatch.setattr(DeepQAPlanNode, "run", fake_plan_run)
-    monkeypatch.setattr(DeepQASearchDispatchNode, "run", fake_search_dispatch_run)
-    monkeypatch.setattr(DeepQADecisionNode, "run", fake_decision_run)
-    monkeypatch.setattr(DeepQASynthesizeNode, "run", fake_synthesize_run)
-
-    start = DeepQAPlanNode(provider="ollama", model="test")
     result = await graph.run(start_node=start, state=state, deps=deps)
 
-    assert "[python.md]" in result.output.answer
+    # Verify citations flag was used
     assert state.context.use_citations is True
+    assert result.output.answer is not None
+    assert isinstance(result.output.sources, list)
+
+    client.close()
 
 
 @pytest.mark.asyncio

tests/test_research_graph_integration.py

@@ -1,32 +1,21 @@
-from typing import Any, cast
-
 import pytest
+from pydantic_ai.models.test import TestModel
 
-from haiku.rag.graph.models import SearchAnswer
+from haiku.rag.client import HaikuRAG
-from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode, DecisionNode
 from haiku.rag.graph.nodes.plan import PlanNode
-from haiku.rag.graph.nodes.search import SearchDispatchNode
-from haiku.rag.graph.nodes.synthesize import SynthesizeNode
 from haiku.rag.research.dependencies import ResearchContext
 from haiku.rag.research.graph import (
     ResearchDeps,
     ResearchState,
     build_research_graph,
 )
-from haiku.rag.research.models import (
+from haiku.rag.research.models import ResearchReport
-    EvaluationResult,
-    GapRecord,
-    GapSeverity,
-    InsightAnalysis,
-    InsightRecord,
-    InsightStatus,
-    ResearchReport,
-)
 from haiku.rag.research.stream import stream_research_graph
 
 
 @pytest.mark.asyncio
-async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
+async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
+    """Test research graph with mocked LLM using TestModel."""
     graph = build_research_graph()
 
     state = ResearchState(
@@ -35,103 +24,43 @@ async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
         confidence_threshold=0.5,
         max_concurrency=2,
     )
-    deps = ResearchDeps(
-        client=cast(Any, None), console=None
-    )  # client unused in patched nodes
 
-    async def fake_plan_run(self, ctx) -> Any:
+    # Use real client but with TestModel for LLM calls
-        ctx.state.context.sub_questions = [
+    client = HaikuRAG(temp_db_path)
-            "Describe haiku.rag in one sentence",
+    deps = ResearchDeps(client=client, console=None)
-            "List core components of haiku.rag",
-        ]
-        ctx.deps.emit_log("planning", ctx.state)
-        return SearchDispatchNode(self.provider, self.model)
 
-    async def fake_search_dispatch_run(self, ctx) -> Any:
+    # Mock get_model to return TestModel which generates valid schema-compliant data
-        # Answer all pending questions deterministically, then move to analysis
+    # Need to patch in all modules that import it
-        while ctx.state.context.sub_questions:
+    def test_model_factory(provider, model):
-            q = ctx.state.context.sub_questions.pop(0)
+        return TestModel()
-            # pydantic BaseModel kwargs not fully typed for pyright
-            ctx.state.context.add_qa_response(
-                SearchAnswer(query=q, answer="A", context=["x"], sources=["s"])  # pyright: ignore[reportCallIssue]
-            )
-            ctx.deps.emit_log(f"answered:{q}", ctx.state)
-        return AnalyzeInsightsNode(self.provider, self.model)
 
-    async def fake_analyze_run(self, ctx) -> Any:
+    monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
-        analysis = InsightAnalysis(
+    monkeypatch.setattr("haiku.rag.graph.nodes.plan.get_model", test_model_factory)
-            highlights=[
+    monkeypatch.setattr("haiku.rag.graph.nodes.search.get_model", test_model_factory)
-                InsightRecord(
+    monkeypatch.setattr("haiku.rag.graph.nodes.analysis.get_model", test_model_factory)
-                    summary="haiku.rag orchestrates research stages",
-                    status=InsightStatus.VALIDATED,
-                    supporting_sources=["s"],
-                    originating_questions=["Describe haiku.rag in one sentence"],
-                )
-            ],
-            gap_assessments=[
-                GapRecord(
-                    description="Need a final summary",
-                    severity=GapSeverity.LOW,
-                    blocking=False,
-                    resolved=False,
-                )
-            ],
-            resolved_gaps=[],
-            new_questions=[],
-            commentary="Insights captured for synthesis",
-        )
-        ctx.state.context.integrate_analysis(analysis)
-        ctx.state.last_analysis = analysis
-        ctx.deps.emit_log("analysis", ctx.state)
-        return DecisionNode(self.provider, self.model)
-
-    async def fake_decision_run(self, ctx) -> Any:
-        ctx.state.last_eval = EvaluationResult(
-            key_insights=["haiku.rag coordinates planning, search, and synthesis"],
-            new_questions=[],
-            gaps=["Need a final summary"],
-            confidence_score=1.0,
-            is_sufficient=True,
-            reasoning="done",
-        )
-        ctx.state.iterations += 1
-        ctx.deps.emit_log("decision", ctx.state)
-        return SynthesizeNode(self.provider, self.model)
-
-    async def fake_synthesize_run(self, ctx) -> Any:
-        report = ResearchReport(
-            title="Haiku RAG",
-            executive_summary="...",
-            main_findings=["f1"],
-            conclusions=["c1"],
-            limitations=[],
-            recommendations=[],
-            sources_summary="s",
-        )
-        from pydantic_graph import End
-
-        return End(report)
-
-    monkeypatch.setattr(PlanNode, "run", fake_plan_run, raising=False)
     monkeypatch.setattr(
-        SearchDispatchNode, "run", fake_search_dispatch_run, raising=False
+        "haiku.rag.graph.nodes.synthesize.get_model", test_model_factory
     )
-    monkeypatch.setattr(AnalyzeInsightsNode, "run", fake_analyze_run, raising=False)
-    monkeypatch.setattr(DecisionNode, "run", fake_decision_run, raising=False)
-    monkeypatch.setattr(SynthesizeNode, "run", fake_synthesize_run, raising=False)
 
     start = PlanNode(provider="test", model="test")
 
     collected = []
+    report = None
     async for event in stream_research_graph(graph, start, state, deps):
         collected.append(event)
         if event.type == "report":
             report = event.report
             break
-    else:  # pragma: no cover - defensive guard
+        elif event.type == "error":
-        report = None
+            pytest.fail(f"Graph execution failed: {event.error}")
 
+    # TestModel will generate valid structured output for each node
+    assert report is not None, (
+        f"No report generated. Events collected: {[e.type for e in collected]}"
+    )
     assert isinstance(report, ResearchReport)
-    assert report.title == "Haiku RAG"
+    assert report.title is not None
-    assert len(state.context.qa_responses) == 2
+    assert isinstance(report.title, str)
     assert any(evt.type == "log" for evt in collected)
+
+    client.close()