haiku.rag/tests/graph/test_research_graph.py

37 lines
1.1 KiB
Python

import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.models import ResearchReport
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
@pytest.mark.vcr()
async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus):
"""Test research graph with real LLM calls recorded via VCR."""
graph = build_research_graph()
client = HaikuRAG(temp_db_path, create=True)
doc = qa_corpus[0]
await client.create_document(
content=doc["document_extracted"], uri=doc["document_id"]
)
state = ResearchState(
context=ResearchContext(original_question=doc["question"]),
max_iterations=1,
confidence_threshold=0.5,
max_concurrency=1,
)
deps = ResearchDeps(client=client)
result = await graph.run(state=state, deps=deps)
assert result is not None
assert isinstance(result, ResearchReport)
assert result.title
assert result.executive_summary
client.close()