41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
import pytest
|
|
from pydantic_ai.models.test import TestModel
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
|
|
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
|
|
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
|
|
"""Test deep Q&A graph with mocked LLM using TestModel."""
|
|
|
|
# Mock get_model to return TestModel which generates valid schema-compliant data
|
|
def test_model_factory(provider, model, config=None):
|
|
return TestModel()
|
|
|
|
# Patch all locations where get_model is imported
|
|
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
|
|
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
|
|
monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory)
|
|
monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
|
|
|
|
graph = build_deep_qa_graph()
|
|
|
|
state = DeepQAState(
|
|
context=DeepQAContext(original_question="What is haiku.rag?"),
|
|
max_sub_questions=3,
|
|
)
|
|
|
|
# Use real client but with TestModel for LLM calls
|
|
client = HaikuRAG(temp_db_path, create=True)
|
|
deps = DeepQADeps(client=client)
|
|
|
|
result = await graph.run(state=state, deps=deps)
|
|
|
|
# TestModel will generate valid structured output based on schemas
|
|
assert result.answer is not None
|
|
assert isinstance(result.answer, str)
|
|
|
|
client.close()
|