89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
from haiku.rag.research.dependencies import ResearchContext
|
|
from haiku.rag.research.graph import (
|
|
EvaluateNode,
|
|
PlanNode,
|
|
ResearchDeps,
|
|
ResearchState,
|
|
SearchDispatchNode,
|
|
SynthesizeNode,
|
|
build_research_graph,
|
|
)
|
|
from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_graph_end_to_end_with_patched_nodes(monkeypatch):
|
|
graph = build_research_graph()
|
|
|
|
state = ResearchState(
|
|
question="What is haiku.rag?",
|
|
context=ResearchContext(original_question="What is haiku.rag?"),
|
|
max_iterations=1,
|
|
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:
|
|
ctx.state.sub_questions = [
|
|
"Describe haiku.rag in one sentence",
|
|
"List core components of haiku.rag",
|
|
]
|
|
return SearchDispatchNode(self.provider, self.model)
|
|
|
|
async def fake_search_dispatch_run(self, ctx) -> Any:
|
|
# Answer all pending questions deterministically, then move to evaluation
|
|
while ctx.state.sub_questions:
|
|
q = ctx.state.sub_questions.pop(0)
|
|
# 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]
|
|
)
|
|
return EvaluateNode(self.provider, self.model)
|
|
|
|
async def fake_evaluate_run(self, ctx) -> Any:
|
|
ctx.state.last_eval = EvaluationResult(
|
|
key_insights=["ok"],
|
|
new_questions=[],
|
|
confidence_score=1.0,
|
|
is_sufficient=True,
|
|
reasoning="done",
|
|
)
|
|
ctx.state.iterations += 1
|
|
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
|
|
)
|
|
monkeypatch.setattr(EvaluateNode, "run", fake_evaluate_run, raising=False)
|
|
monkeypatch.setattr(SynthesizeNode, "run", fake_synthesize_run, raising=False)
|
|
|
|
start = PlanNode(provider="test", model="test")
|
|
|
|
result = await graph.run(start, state=state, deps=deps)
|
|
report = result.output
|
|
|
|
assert isinstance(report, ResearchReport)
|
|
assert report.title == "Haiku RAG"
|
|
assert len(state.context.qa_responses) == 2
|