danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor research, deep ask to work with new beta pyndtic AI graph

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-03 11:49:31 +02:00
parent 8da62e3868
commit bc71ad6fb0
No known key found for this signature in database
15 changed files with 785 additions and 882 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
haiku_rag_slim/haiku/rag/app.py
View file

@ -13,12 +13,8 @@ from haiku.rag.config import Config
from haiku.rag.mcp import create_mcp_server
from haiku.rag.monitor import FileWatcher
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import (
PlanNode,
ResearchDeps,
ResearchState,
build_research_graph,
)
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.state import ResearchDeps, ResearchState
from haiku.rag.research.stream import stream_research_graph
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
@ -215,10 +211,12 @@ class HaikuRAGApp:
from haiku.rag.qa.deep.dependencies import DeepQAContext
from haiku.rag.qa.deep.graph import build_deep_qa_graph
from haiku.rag.qa.deep.nodes import DeepQAPlanNode
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
graph = build_deep_qa_graph()
graph = build_deep_qa_graph(
provider=Config.qa.provider,
model=Config.qa.model,
)
context = DeepQAContext(
original_question=question, use_citations=cite
)
@ -227,15 +225,8 @@ class HaikuRAGApp:
client=self.client, console=Console() if verbose else None
)
start_node = DeepQAPlanNode(
provider=Config.qa.provider,
model=Config.qa.model,
)
result = await graph.run(
start_node=start_node, state=state, deps=deps
)
answer = result.output.answer
result = await graph.run(state=state, deps=deps)
answer = result.answer
else:
answer = await self.client.ask(question, cite=cite)
@ -262,7 +253,10 @@ class HaikuRAGApp:
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
graph = build_research_graph()
graph = build_research_graph(
provider=Config.research.provider or Config.qa.provider,
model=Config.research.model or Config.qa.model,
)
context = ResearchContext(original_question=question)
state = ResearchState(
context=context,
@ -274,12 +268,8 @@ class HaikuRAGApp:
client=client, console=self.console if verbose else None
)
start = PlanNode(
provider=Config.research.provider or Config.qa.provider,
model=Config.research.model or Config.qa.model,
)
report = None
async for event in stream_research_graph(graph, start, state, deps):
async for event in stream_research_graph(graph, state, deps):
if event.type == "report":
report = event.report
break

1
haiku_rag_slim/haiku/rag/graph/__init__.py
View file

@ -1 +0,0 @@
from haiku.rag.graph.models import ResearchPlan, SearchAnswer

31
haiku_rag_slim/haiku/rag/graph/base.py
View file

@ -1,31 +0,0 @@
from typing import Protocol, runtime_checkable
from pydantic import BaseModel, Field
from rich.console import Console
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.models import SearchAnswer
@runtime_checkable
class GraphContext(Protocol):
"""Protocol for graph context objects."""
original_question: str
sub_questions: list[str]
qa_responses: list[SearchAnswer]
def add_qa_response(self, qa: SearchAnswer) -> None: ...
class BaseGraphDeps(BaseModel):
"""Base dependencies for graph nodes."""
model_config = {"arbitrary_types_allowed": True}
client: HaikuRAG = Field(description="RAG client for document operations")
console: Console | None = None
def emit_log(self, message: str) -> None:
if self.console:
self.console.print(message)

182
haiku_rag_slim/haiku/rag/graph/nodes/analysis.py
View file

@ -1,182 +0,0 @@
from dataclasses import dataclass
from pydantic_ai import Agent
from pydantic_graph import BaseNode, GraphRunContext
from haiku.rag.graph.common import get_model, log
from haiku.rag.research.common import (
format_analysis_for_prompt,
format_context_for_prompt,
)
from haiku.rag.research.dependencies import ResearchDependencies
from haiku.rag.research.models import EvaluationResult, InsightAnalysis, ResearchReport
from haiku.rag.research.prompts import DECISION_AGENT_PROMPT, INSIGHT_AGENT_PROMPT
from haiku.rag.research.state import ResearchDeps, ResearchState
@dataclass
class AnalyzeInsightsNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
provider: str
model: str
async def run(
self, ctx: GraphRunContext[ResearchState, ResearchDeps]
) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]🧭 Synthesizing new insights and gap status...[/bold cyan]",
)
agent = Agent(
model=get_model(self.provider, self.model),
output_type=InsightAnalysis,
instructions=INSIGHT_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Review the latest research context and update the shared ledger of insights, gaps,"
" and follow-up questions.\n\n"
f"{context_xml}"
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
result = await agent.run(prompt, deps=agent_deps)
analysis: InsightAnalysis = result.output
state.context.integrate_analysis(analysis)
state.last_analysis = analysis
if analysis.commentary:
log(deps, state, f" Summary: {analysis.commentary}")
if analysis.highlights:
log(deps, state, " [bold]Updated insights:[/bold]")
for insight in analysis.highlights:
label = insight.status.value
log(
deps,
state,
f" • ({label}) {insight.summary}",
)
if analysis.gap_assessments:
log(deps, state, " [bold yellow]Gap updates:[/bold yellow]")
for gap in analysis.gap_assessments:
status = "resolved" if gap.resolved else "open"
severity = gap.severity.value
log(
deps,
state,
f" • ({severity}/{status}) {gap.description}",
)
if analysis.resolved_gaps:
log(deps, state, " [green]Resolved gaps:[/green]")
for resolved in analysis.resolved_gaps:
log(deps, state, f" • {resolved}")
if analysis.new_questions:
log(deps, state, " [cyan]Proposed follow-ups:[/cyan]")
for question in analysis.new_questions:
log(deps, state, f" • {question}")
from haiku.rag.graph.nodes.analysis import DecisionNode
return DecisionNode(self.provider, self.model)
@dataclass
class DecisionNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
provider: str
model: str
async def run(
self, ctx: GraphRunContext[ResearchState, ResearchDeps]
) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📊 Evaluating research sufficiency...[/bold cyan]",
)
agent = Agent(
model=get_model(self.provider, self.model),
output_type=EvaluationResult,
instructions=DECISION_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
analysis_xml = format_analysis_for_prompt(state.last_analysis)
prompt_parts = [
"Assess whether the research now answers the original question with adequate confidence.",
context_xml,
analysis_xml,
]
if state.last_eval is not None:
prev = state.last_eval
prompt_parts.append(
"<previous_evaluation>"
f"<confidence>{prev.confidence_score:.2f}</confidence>"
f"<is_sufficient>{str(prev.is_sufficient).lower()}</is_sufficient>"
f"<reasoning>{prev.reasoning}</reasoning>"
"</previous_evaluation>"
)
prompt = "\n\n".join(part for part in prompt_parts if part)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
decision_result = await agent.run(prompt, deps=agent_deps)
output = decision_result.output
state.last_eval = output
state.iterations += 1
for new_q in output.new_questions:
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
if output.key_insights:
log(deps, state, " [bold]Key insights:[/bold]")
for insight in output.key_insights:
log(deps, state, f" • {insight}")
if output.gaps:
log(deps, state, " [bold yellow]Remaining gaps:[/bold yellow]")
for gap in output.gaps:
log(deps, state, f" • {gap}")
log(
deps,
state,
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
)
status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
log(deps, state, f" Sufficient: {status}")
from haiku.rag.graph.nodes.search import SearchDispatchNode
from haiku.rag.graph.nodes.synthesize import SynthesizeNode
if (
output.is_sufficient
and output.confidence_score >= state.confidence_threshold
) or state.iterations >= state.max_iterations:
log(deps, state, "\n[bold green]✅ Stopping research.[/bold green]")
return SynthesizeNode(self.provider, self.model)
return SearchDispatchNode(self.provider, self.model)

72
haiku_rag_slim/haiku/rag/graph/nodes/plan.py
View file

@ -1,72 +0,0 @@
from dataclasses import dataclass
from pydantic_ai import Agent, RunContext
from pydantic_graph import BaseNode, GraphRunContext
from haiku.rag.graph.common import get_model, log
from haiku.rag.graph.models import ResearchPlan
from haiku.rag.graph.prompts import PLAN_PROMPT
from haiku.rag.research.dependencies import ResearchDependencies
from haiku.rag.research.models import ResearchReport
from haiku.rag.research.state import ResearchDeps, ResearchState
@dataclass
class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
provider: str
model: str
async def run(
self, ctx: GraphRunContext[ResearchState, ResearchDeps]
) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
state = ctx.state
deps = ctx.deps
log(deps, state, "\n[bold cyan]📋 Creating research plan...[/bold cyan]")
plan_agent = Agent(
model=get_model(self.provider, self.model),
output_type=ResearchPlan,
instructions=(
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
retries=3,
deps_type=ResearchDependencies,
)
@plan_agent.tool
async def gather_context(
ctx2: RunContext[ResearchDependencies], query: str, limit: int = 6
) -> str:
results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(results)
return "\n\n".join(chunk.content for chunk, _ in expanded)
prompt = (
"Plan a focused research approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)
log(deps, state, "\n[bold green]✅ Research Plan Created:[/bold green]")
log(
deps,
state,
f" [bold]Main Question:[/bold] {state.context.original_question}",
)
log(deps, state, " [bold]Sub-questions:[/bold]")
for i, sq in enumerate(state.context.sub_questions, 1):
log(deps, state, f" {i}. {sq}")
from haiku.rag.graph.nodes.search import SearchDispatchNode
return SearchDispatchNode(self.provider, self.model)

97
haiku_rag_slim/haiku/rag/graph/nodes/search.py
View file

@ -1,97 +0,0 @@
import asyncio
from dataclasses import dataclass
from typing import Any
from pydantic_ai import Agent, RunContext
from pydantic_ai.format_prompt import format_as_xml
from pydantic_ai.output import ToolOutput
from pydantic_graph import BaseNode, GraphRunContext
from haiku.rag.graph.common import get_model, log
from haiku.rag.graph.models import SearchAnswer
from haiku.rag.graph.prompts import SEARCH_AGENT_PROMPT
from haiku.rag.research.dependencies import ResearchDependencies
from haiku.rag.research.models import ResearchReport
from haiku.rag.research.state import ResearchDeps, ResearchState
@dataclass
class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
provider: str
model: str
async def run(
self, ctx: GraphRunContext[ResearchState, ResearchDeps]
) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
state = ctx.state
deps = ctx.deps
if not state.context.sub_questions:
from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode
return AnalyzeInsightsNode(self.provider, self.model)
# Take up to max_concurrency questions and answer them concurrently
take = max(1, state.max_concurrency)
batch: list[str] = []
while state.context.sub_questions and len(batch) < take:
batch.append(state.context.sub_questions.pop(0))
async def answer_one(sub_q: str) -> SearchAnswer | None:
log(
deps,
state,
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
)
agent = Agent(
model=get_model(self.provider, self.model),
output_type=ToolOutput(SearchAnswer, max_retries=3),
instructions=SEARCH_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
@agent.tool
async def search_and_answer(
ctx2: RunContext[ResearchDependencies], query: str, limit: int = 5
) -> str:
search_results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(search_results)
entries: list[dict[str, Any]] = [
{
"text": chunk.content,
"score": score,
"document_uri": (
chunk.document_title or chunk.document_uri or ""
),
}
for chunk, score in expanded
]
if not entries:
return f"No relevant information found in the knowledge base for: {query}"
return format_as_xml(entries, root_tag="snippets")
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
try:
result = await agent.run(sub_q, deps=agent_deps)
except Exception as e:
log(deps, state, f"[red]Search failed:[/red] {e}")
return None
return result.output
answers = await asyncio.gather(*(answer_one(q) for q in batch))
for ans in answers:
if ans is None:
continue
state.context.add_qa_response(ans)
preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
log(deps, state, f" [green]✓[/green] {preview}")
return SearchDispatchNode(self.provider, self.model)

54
haiku_rag_slim/haiku/rag/graph/nodes/synthesize.py
View file

@ -1,54 +0,0 @@
from dataclasses import dataclass
from pydantic_ai import Agent
from pydantic_graph import BaseNode, End, GraphRunContext
from haiku.rag.graph.common import get_model, log
from haiku.rag.research.common import format_context_for_prompt
from haiku.rag.research.dependencies import ResearchDependencies
from haiku.rag.research.models import ResearchReport
from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT
from haiku.rag.research.state import ResearchDeps, ResearchState
@dataclass
class SynthesizeNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
provider: str
model: str
async def run(
self, ctx: GraphRunContext[ResearchState, ResearchDeps]
) -> End[ResearchReport]:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📝 Generating final research report...[/bold cyan]",
)
agent = Agent(
model=get_model(self.provider, self.model),
output_type=ResearchReport,
instructions=SYNTHESIS_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Generate a comprehensive research report based on all gathered information.\n\n"
f"{context_xml}\n\n"
"Create a detailed report that synthesizes all findings into a coherent response."
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
result = await agent.run(prompt, deps=agent_deps)
log(deps, state, "[bold green]✅ Research complete![/bold green]")
return End(result.output)

34
haiku_rag_slim/haiku/rag/mcp.py
View file

@ -194,25 +194,20 @@ def create_mcp_server(db_path: Path) -> FastMCP:
from haiku.rag.config import Config
from haiku.rag.qa.deep.dependencies import DeepQAContext
from haiku.rag.qa.deep.graph import build_deep_qa_graph
from haiku.rag.qa.deep.nodes import DeepQAPlanNode
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
graph = build_deep_qa_graph()
graph = build_deep_qa_graph(
provider=Config.qa.provider,
model=Config.qa.model,
)
context = DeepQAContext(
original_question=question, use_citations=cite
)
state = DeepQAState(context=context)
deps = DeepQADeps(client=rag)
start_node = DeepQAPlanNode(
provider=Config.qa.provider,
model=Config.qa.model,
)
result = await graph.run(
start_node=start_node, state=state, deps=deps
)
answer = result.output.answer
result = await graph.run(state=state, deps=deps)
answer = result.answer
else:
answer = await rag.ask(question, cite=cite)
return answer
@ -241,13 +236,15 @@ def create_mcp_server(db_path: Path) -> FastMCP:
A research report with findings, or None if an error occurred.
"""
try:
from haiku.rag.graph.nodes.plan import PlanNode
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.state import ResearchDeps, ResearchState
async with HaikuRAG(db_path) as rag:
graph = build_research_graph()
graph = build_research_graph(
provider=Config.research.provider or Config.qa.provider,
model=Config.research.model or Config.qa.model,
)
state = ResearchState(
context=ResearchContext(original_question=question),
max_iterations=max_iterations,
@ -256,16 +253,9 @@ def create_mcp_server(db_path: Path) -> FastMCP:
)
deps = ResearchDeps(client=rag)
result = await graph.run(
PlanNode(
provider=Config.research.provider or Config.qa.provider,
model=Config.research.model or Config.qa.model,
),
state=state,
deps=deps,
)
result = await graph.run(state=state, deps=deps)
return result.output
return result
except Exception:
return None

341
haiku_rag_slim/haiku/rag/qa/deep/graph.py
View file

@ -1,21 +1,332 @@
from pydantic_graph import Graph
from typing import Any
from haiku.rag.qa.deep.models import DeepQAAnswer
from haiku.rag.qa.deep.nodes import (
DeepQADecisionNode,
DeepQAPlanNode,
DeepQASearchDispatchNode,
DeepQASynthesizeNode,
from pydantic_ai import Agent, RunContext
from pydantic_ai.format_prompt import format_as_xml
from pydantic_ai.output import ToolOutput
from pydantic_graph.beta import GraphBuilder, StepContext
from haiku.rag.graph.common import get_model, log
from haiku.rag.graph.models import ResearchPlan, SearchAnswer
from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
from haiku.rag.qa.deep.dependencies import DeepQADependencies
from haiku.rag.qa.deep.models import DeepQAAnswer, DeepQAEvaluation
from haiku.rag.qa.deep.prompts import (
DECISION_PROMPT,
SYNTHESIS_PROMPT,
SYNTHESIS_PROMPT_WITH_CITATIONS,
)
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
def build_deep_qa_graph() -> Graph[DeepQAState, DeepQADeps, DeepQAAnswer]:
return Graph(
nodes=[
DeepQAPlanNode,
DeepQASearchDispatchNode,
DeepQADecisionNode,
DeepQASynthesizeNode,
]
def build_deep_qa_graph(provider: str, model: str):
g = GraphBuilder(
state_type=DeepQAState,
deps_type=DeepQADeps,
output_type=DeepQAAnswer,
)
@g.step
async def plan(ctx: StepContext[DeepQAState, DeepQADeps, None]) -> None:
state = ctx.state
deps = ctx.deps
log(deps, state, "\n[bold cyan]📋 Planning approach...[/bold cyan]")
plan_agent = Agent(
model=get_model(provider, model),
output_type=ResearchPlan,
instructions=(
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
retries=3,
deps_type=DeepQADependencies,
)
@plan_agent.tool
async def gather_context(
ctx2: RunContext[DeepQADependencies], query: str, limit: int = 6
) -> str:
results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(results)
return "\n\n".join(chunk.content for chunk, _ in expanded)
prompt = (
"Plan a focused approach for answering the main question.\n\n"
f"Main question: {state.context.original_question}"
)
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)[
: state.max_sub_questions
]
log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]")
log(
deps,
state,
f" [bold]Main Question:[/bold] {state.context.original_question}",
)
log(deps, state, " [bold]Sub-questions:[/bold]")
for i, sq in enumerate(state.context.sub_questions, 1):
log(deps, state, f" {i}. {sq}")
@g.step
async def search_one(
ctx: StepContext[DeepQAState, DeepQADeps, str],
) -> SearchAnswer | None:
state = ctx.state
deps = ctx.deps
sub_q = ctx.inputs
log(
deps,
state,
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
)
agent = Agent(
model=get_model(provider, model),
output_type=ToolOutput(SearchAnswer, max_retries=3),
instructions=SEARCH_AGENT_PROMPT,
retries=3,
deps_type=DeepQADependencies,
)
@agent.tool
async def search_and_answer(
ctx2: RunContext[DeepQADependencies], query: str, limit: int = 5
) -> str:
search_results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(search_results)
entries: list[dict[str, Any]] = [
{
"text": chunk.content,
"score": score,
"document_uri": (chunk.document_title or chunk.document_uri or ""),
}
for chunk, score in expanded
]
if not entries:
return (
f"No relevant information found in the knowledge base for: {query}"
)
return format_as_xml(entries, root_tag="snippets")
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
try:
result = await agent.run(sub_q, deps=agent_deps)
except Exception as e:
log(deps, state, f"[red]Search failed:[/red] {e}")
return None
answer = result.output
if answer:
state.context.add_qa_response(answer)
preview = answer.answer[:150] + ("…" if len(answer.answer) > 150 else "")
log(deps, state, f" [green]✓[/green] {preview}")
return answer
@g.step
async def decide(
ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer | None]],
) -> bool:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📊 Evaluating information sufficiency...[/bold cyan]",
)
agent = Agent(
model=get_model(provider, model),
output_type=DeepQAEvaluation,
instructions=DECISION_PROMPT,
retries=3,
deps_type=DeepQADependencies,
)
context_data = {
"original_question": state.context.original_question,
"gathered_answers": [
{
"question": qa.query,
"answer": qa.answer,
"sources": qa.sources,
}
for qa in state.context.qa_responses
],
}
context_xml = format_as_xml(context_data, root_tag="gathered_information")
prompt = (
"Evaluate whether we have sufficient information to answer the question.\n\n"
f"{context_xml}"
)
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
result = await agent.run(prompt, deps=agent_deps)
evaluation = result.output
state.iterations += 1
log(deps, state, f" [bold]Assessment:[/bold] {evaluation.reasoning}")
status = "[green]Yes[/green]" if evaluation.is_sufficient else "[red]No[/red]"
log(deps, state, f" Sufficient: {status}")
for new_q in evaluation.new_questions:
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
if evaluation.new_questions:
log(deps, state, " [cyan]New questions:[/cyan]")
for question in evaluation.new_questions:
log(deps, state, f" • {question}")
should_continue = (
not evaluation.is_sufficient and state.iterations < state.max_iterations
)
if not should_continue:
if state.iterations >= state.max_iterations:
log(
deps,
state,
f"\n[bold yellow]⚠️ Reached max iterations ({state.max_iterations})[/bold yellow]",
)
log(deps, state, "\n[bold green]✅ Moving to synthesis.[/bold green]")
else:
log(
deps,
state,
f"\n[bold cyan]🔄 Starting iteration {state.iterations + 1}...[/bold cyan]",
)
return should_continue
@g.step
async def get_batch(
ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
) -> list[str] | None:
"""Get next batch of questions from state."""
state = ctx.state
take = max(1, state.max_concurrency)
batch: list[str] = []
while state.context.sub_questions and len(batch) < take:
batch.append(state.context.sub_questions.pop(0))
return batch if batch else None
@g.step
async def synthesize(
ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
) -> DeepQAAnswer:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📝 Synthesizing final answer...[/bold cyan]",
)
prompt_template = (
SYNTHESIS_PROMPT_WITH_CITATIONS
if state.context.use_citations
else SYNTHESIS_PROMPT
)
agent = Agent(
model=get_model(provider, model),
output_type=DeepQAAnswer,
instructions=prompt_template,
retries=3,
deps_type=DeepQADependencies,
)
context_data = {
"original_question": state.context.original_question,
"sub_answers": [
{
"question": qa.query,
"answer": qa.answer,
"sources": qa.sources,
}
for qa in state.context.qa_responses
],
}
context_xml = format_as_xml(context_data, root_tag="gathered_information")
prompt = f"Synthesize a comprehensive answer to the original question.\n\n{context_xml}"
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
result = await agent.run(prompt, deps=agent_deps)
log(deps, state, "[bold green]✅ Answer complete![/bold green]")
return result.output
# Build the graph structure
def collect_reducer(
acc: list[SearchAnswer | None], item: SearchAnswer | None
) -> list[SearchAnswer | None]:
return acc + [item] if item else acc
collect_answers = g.join(
collect_reducer,
initial_factory=lambda: [],
)
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
)
# Branch based on whether we have questions
g.add(
g.edge_from(get_batch).to(
g.decision()
.branch(g.match(list).label("Has questions").map().to(search_one))
.branch(g.match(type(None)).label("No questions").to(synthesize))
),
g.edge_from(search_one).to(collect_answers),
g.edge_from(collect_answers).to(decide),
)
# Branch based on decision
g.add(
g.edge_from(decide).to(
g.decision()
.branch(
g.match(bool, matches=lambda x: x).label("Continue QA").to(get_batch)
)
.branch(
g.match(bool, matches=lambda x: not x)
.label("Done with QA")
.to(synthesize)
)
),
g.edge_from(synthesize).to(g.end_node),
)
return g.build()

303
haiku_rag_slim/haiku/rag/qa/deep/nodes.py
View file

@ -1,303 +0,0 @@
import asyncio
from dataclasses import dataclass
from typing import Any
from pydantic_ai import Agent, RunContext
from pydantic_ai.format_prompt import format_as_xml
from pydantic_ai.output import ToolOutput
from pydantic_graph import BaseNode, End, GraphRunContext
from haiku.rag.graph.common import get_model, log
from haiku.rag.graph.models import ResearchPlan, SearchAnswer
from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
from haiku.rag.qa.deep.dependencies import DeepQADependencies
from haiku.rag.qa.deep.models import DeepQAAnswer, DeepQAEvaluation
from haiku.rag.qa.deep.prompts import (
DECISION_PROMPT,
SYNTHESIS_PROMPT,
SYNTHESIS_PROMPT_WITH_CITATIONS,
)
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
@dataclass
class DeepQAPlanNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
provider: str
model: str
async def run(
self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
state = ctx.state
deps = ctx.deps
log(deps, state, "\n[bold cyan]📋 Planning approach...[/bold cyan]")
plan_agent = Agent(
model=get_model(self.provider, self.model),
output_type=ResearchPlan,
instructions=(
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
retries=3,
deps_type=DeepQADependencies,
)
@plan_agent.tool
async def gather_context(
ctx2: RunContext[DeepQADependencies], query: str, limit: int = 6
) -> str:
results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(results)
return "\n\n".join(chunk.content for chunk, _ in expanded)
prompt = (
"Plan a focused approach for answering the main question.\n\n"
f"Main question: {state.context.original_question}"
)
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)[
: state.max_sub_questions
]
log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]")
log(
deps,
state,
f" [bold]Main Question:[/bold] {state.context.original_question}",
)
log(deps, state, " [bold]Sub-questions:[/bold]")
for i, sq in enumerate(state.context.sub_questions, 1):
log(deps, state, f" {i}. {sq}")
return DeepQASearchDispatchNode(self.provider, self.model)
@dataclass
class DeepQASearchDispatchNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
provider: str
model: str
async def run(
self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
state = ctx.state
deps = ctx.deps
if not state.context.sub_questions:
return DeepQADecisionNode(self.provider, self.model)
# Take up to max_concurrency questions and answer them concurrently
take = max(1, state.max_concurrency)
batch: list[str] = []
while state.context.sub_questions and len(batch) < take:
batch.append(state.context.sub_questions.pop(0))
async def answer_one(sub_q: str) -> SearchAnswer | None:
log(
deps,
state,
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
)
agent = Agent(
model=get_model(self.provider, self.model),
output_type=ToolOutput(SearchAnswer, max_retries=3),
instructions=SEARCH_AGENT_PROMPT,
retries=3,
deps_type=DeepQADependencies,
)
@agent.tool
async def search_and_answer(
ctx2: RunContext[DeepQADependencies], query: str, limit: int = 5
) -> str:
search_results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(search_results)
entries: list[dict[str, Any]] = [
{
"text": chunk.content,
"score": score,
"document_uri": (
chunk.document_title or chunk.document_uri or ""
),
}
for chunk, score in expanded
]
if not entries:
return f"No relevant information found in the knowledge base for: {query}"
return format_as_xml(entries, root_tag="snippets")
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
try:
result = await agent.run(sub_q, deps=agent_deps)
except Exception as e:
log(deps, state, f"[red]Search failed:[/red] {e}")
return None
return result.output
answers = await asyncio.gather(*(answer_one(q) for q in batch))
for ans in answers:
if ans is None:
continue
state.context.add_qa_response(ans)
preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
log(deps, state, f" [green]✓[/green] {preview}")
return DeepQASearchDispatchNode(self.provider, self.model)
@dataclass
class DeepQADecisionNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
provider: str
model: str
async def run(
self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📊 Evaluating information sufficiency...[/bold cyan]",
)
agent = Agent(
model=get_model(self.provider, self.model),
output_type=DeepQAEvaluation,
instructions=DECISION_PROMPT,
retries=3,
deps_type=DeepQADependencies,
)
context_data = {
"original_question": state.context.original_question,
"gathered_answers": [
{
"question": qa.query,
"answer": qa.answer,
"sources": qa.sources,
}
for qa in state.context.qa_responses
],
}
context_xml = format_as_xml(context_data, root_tag="gathered_information")
prompt = (
"Evaluate whether we have sufficient information to answer the question.\n\n"
f"{context_xml}"
)
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
result = await agent.run(prompt, deps=agent_deps)
evaluation = result.output
state.iterations += 1
log(deps, state, f" [bold]Assessment:[/bold] {evaluation.reasoning}")
status = "[green]Yes[/green]" if evaluation.is_sufficient else "[red]No[/red]"
log(deps, state, f" Sufficient: {status}")
# Add new questions if not sufficient
for new_q in evaluation.new_questions:
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
if evaluation.new_questions:
log(deps, state, " [cyan]New questions:[/cyan]")
for question in evaluation.new_questions:
log(deps, state, f" • {question}")
# Decide next step
if evaluation.is_sufficient or state.iterations >= state.max_iterations:
if state.iterations >= state.max_iterations:
log(
deps,
state,
f"\n[bold yellow]⚠️ Reached max iterations ({state.max_iterations})[/bold yellow]",
)
log(deps, state, "\n[bold green]✅ Moving to synthesis.[/bold green]")
return DeepQASynthesizeNode(self.provider, self.model)
log(
deps,
state,
f"\n[bold cyan]🔄 Starting iteration {state.iterations + 1}...[/bold cyan]",
)
return DeepQASearchDispatchNode(self.provider, self.model)
@dataclass
class DeepQASynthesizeNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
provider: str
model: str
async def run(
self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
) -> End[DeepQAAnswer]:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📝 Synthesizing final answer...[/bold cyan]",
)
prompt_template = (
SYNTHESIS_PROMPT_WITH_CITATIONS
if state.context.use_citations
else SYNTHESIS_PROMPT
)
agent = Agent(
model=get_model(self.provider, self.model),
output_type=DeepQAAnswer,
instructions=prompt_template,
retries=3,
deps_type=DeepQADependencies,
)
context_data = {
"original_question": state.context.original_question,
"sub_answers": [
{
"question": qa.query,
"answer": qa.answer,
"sources": qa.sources,
}
for qa in state.context.qa_responses
],
}
context_xml = format_as_xml(context_data, root_tag="gathered_information")
prompt = f"Synthesize a comprehensive answer to the original question.\n\n{context_xml}"
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
result = await agent.run(prompt, deps=agent_deps)
log(deps, state, "[bold green]✅ Answer complete![/bold green]")
return End(result.output)

406
haiku_rag_slim/haiku/rag/research/graph.py
View file

@ -1,20 +1,396 @@
from pydantic_graph import Graph
from typing import Any
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.models import ResearchReport
from pydantic_ai import Agent, RunContext
from pydantic_ai.format_prompt import format_as_xml
from pydantic_ai.output import ToolOutput
from pydantic_graph.beta import GraphBuilder, StepContext
from haiku.rag.graph.common import get_model, log
from haiku.rag.graph.models import ResearchPlan, SearchAnswer
from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
from haiku.rag.research.common import (
format_analysis_for_prompt,
format_context_for_prompt,
)
from haiku.rag.research.dependencies import ResearchDependencies
from haiku.rag.research.models import (
EvaluationResult,
InsightAnalysis,
ResearchReport,
)
from haiku.rag.research.prompts import (
DECISION_AGENT_PROMPT,
INSIGHT_AGENT_PROMPT,
SYNTHESIS_AGENT_PROMPT,
)
from haiku.rag.research.state import ResearchDeps, ResearchState
def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]:
return Graph(
nodes=[
PlanNode,
SearchDispatchNode,
AnalyzeInsightsNode,
DecisionNode,
SynthesizeNode,
]
def build_research_graph(provider: str, model: str):
g = GraphBuilder(
state_type=ResearchState,
deps_type=ResearchDeps,
output_type=ResearchReport,
)
@g.step
async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
state = ctx.state
deps = ctx.deps
log(deps, state, "\n[bold cyan]📋 Creating research plan...[/bold cyan]")
plan_agent = Agent(
model=get_model(provider, model),
output_type=ResearchPlan,
instructions=(
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
retries=3,
deps_type=ResearchDependencies,
)
@plan_agent.tool
async def gather_context(
ctx2: RunContext[ResearchDependencies], query: str, limit: int = 6
) -> str:
results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(results)
return "\n\n".join(chunk.content for chunk, _ in expanded)
prompt = (
"Plan a focused research approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)
log(deps, state, "\n[bold green]✅ Research Plan Created:[/bold green]")
log(
deps,
state,
f" [bold]Main Question:[/bold] {state.context.original_question}",
)
log(deps, state, " [bold]Sub-questions:[/bold]")
for i, sq in enumerate(state.context.sub_questions, 1):
log(deps, state, f" {i}. {sq}")
@g.step
async def search_one(
ctx: StepContext[ResearchState, ResearchDeps, str],
) -> SearchAnswer | None:
state = ctx.state
deps = ctx.deps
sub_q = ctx.inputs
log(
deps,
state,
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
)
agent = Agent(
model=get_model(provider, model),
output_type=ToolOutput(SearchAnswer, max_retries=3),
instructions=SEARCH_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
@agent.tool
async def search_and_answer(
ctx2: RunContext[ResearchDependencies], query: str, limit: int = 5
) -> str:
search_results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(search_results)
entries: list[dict[str, Any]] = [
{
"text": chunk.content,
"score": score,
"document_uri": (chunk.document_title or chunk.document_uri or ""),
}
for chunk, score in expanded
]
if not entries:
return (
f"No relevant information found in the knowledge base for: {query}"
)
return format_as_xml(entries, root_tag="snippets")
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
try:
result = await agent.run(sub_q, deps=agent_deps)
except Exception as e:
log(deps, state, f"[red]Search failed:[/red] {e}")
return None
answer = result.output
if answer:
state.context.add_qa_response(answer)
preview = answer.answer[:150] + ("…" if len(answer.answer) > 150 else "")
log(deps, state, f" [green]✓[/green] {preview}")
return answer
@g.step
async def analyze_insights(
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer | None]],
) -> None:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]🧭 Synthesizing new insights and gap status...[/bold cyan]",
)
agent = Agent(
model=get_model(provider, model),
output_type=InsightAnalysis,
instructions=INSIGHT_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Review the latest research context and update the shared ledger of insights, gaps,"
" and follow-up questions.\n\n"
f"{context_xml}"
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
result = await agent.run(prompt, deps=agent_deps)
analysis: InsightAnalysis = result.output
state.context.integrate_analysis(analysis)
state.last_analysis = analysis
if analysis.commentary:
log(deps, state, f" Summary: {analysis.commentary}")
if analysis.highlights:
log(deps, state, " [bold]Updated insights:[/bold]")
for insight in analysis.highlights:
label = insight.status.value
log(
deps,
state,
f" • ({label}) {insight.summary}",
)
if analysis.gap_assessments:
log(deps, state, " [bold yellow]Gap updates:[/bold yellow]")
for gap in analysis.gap_assessments:
status = "resolved" if gap.resolved else "open"
severity = gap.severity.value
log(
deps,
state,
f" • ({severity}/{status}) {gap.description}",
)
if analysis.resolved_gaps:
log(deps, state, " [green]Resolved gaps:[/green]")
for resolved in analysis.resolved_gaps:
log(deps, state, f" • {resolved}")
if analysis.new_questions:
log(deps, state, " [cyan]Proposed follow-ups:[/cyan]")
for question in analysis.new_questions:
log(deps, state, f" • {question}")
@g.step
async def decide(ctx: StepContext[ResearchState, ResearchDeps, None]) -> bool:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📊 Evaluating research sufficiency...[/bold cyan]",
)
agent = Agent(
model=get_model(provider, model),
output_type=EvaluationResult,
instructions=DECISION_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
analysis_xml = format_analysis_for_prompt(state.last_analysis)
prompt_parts = [
"Assess whether the research now answers the original question with adequate confidence.",
context_xml,
analysis_xml,
]
if state.last_eval is not None:
prev = state.last_eval
prompt_parts.append(
"<previous_evaluation>"
f"<confidence>{prev.confidence_score:.2f}</confidence>"
f"<is_sufficient>{str(prev.is_sufficient).lower()}</is_sufficient>"
f"<reasoning>{prev.reasoning}</reasoning>"
"</previous_evaluation>"
)
prompt = "\n\n".join(part for part in prompt_parts if part)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
decision_result = await agent.run(prompt, deps=agent_deps)
output = decision_result.output
state.last_eval = output
state.iterations += 1
for new_q in output.new_questions:
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
if output.key_insights:
log(deps, state, " [bold]Key insights:[/bold]")
for insight in output.key_insights:
log(deps, state, f" • {insight}")
if output.gaps:
log(deps, state, " [bold yellow]Remaining gaps:[/bold yellow]")
for gap in output.gaps:
log(deps, state, f" • {gap}")
log(
deps,
state,
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
)
status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
log(deps, state, f" Sufficient: {status}")
should_continue = (
not output.is_sufficient
or output.confidence_score < state.confidence_threshold
) and state.iterations < state.max_iterations
if not should_continue:
log(deps, state, "\n[bold green]✅ Stopping research.[/bold green]")
return should_continue
@g.step
async def get_batch(
ctx: StepContext[ResearchState, ResearchDeps, None | bool],
) -> list[str] | None:
"""Get next batch of questions from state."""
state = ctx.state
take = max(1, state.max_concurrency)
batch: list[str] = []
while state.context.sub_questions and len(batch) < take:
batch.append(state.context.sub_questions.pop(0))
return batch if batch else None
@g.step
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, None | bool],
) -> ResearchReport:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📝 Generating final research report...[/bold cyan]",
)
agent = Agent(
model=get_model(provider, model),
output_type=ResearchReport,
instructions=SYNTHESIS_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Generate a comprehensive research report based on all gathered information.\n\n"
f"{context_xml}\n\n"
"Create a detailed report that synthesizes all findings into a coherent response."
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
result = await agent.run(prompt, deps=agent_deps)
log(deps, state, "[bold green]✅ Research complete![/bold green]")
return result.output
# Build the graph structure
def collect_reducer(
acc: list[SearchAnswer | None], item: SearchAnswer | None
) -> list[SearchAnswer | None]:
return acc + [item] if item else acc
collect_answers = g.join(
collect_reducer,
initial_factory=lambda: [],
)
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
)
# Branch based on whether we have questions
g.add(
g.edge_from(get_batch).to(
g.decision()
.branch(g.match(list).label("Has questions").map().to(search_one))
.branch(g.match(type(None)).label("No questions").to(synthesize))
),
g.edge_from(search_one).to(collect_answers),
g.edge_from(collect_answers).to(analyze_insights),
g.edge_from(analyze_insights).to(decide),
)
# Branch based on decision
g.add(
g.edge_from(decide).to(
g.decision()
.branch(
g.match(bool, matches=lambda x: x)
.label("Continue research")
.to(get_batch)
)
.branch(
g.match(bool, matches=lambda x: not x)
.label("Done researching")
.to(synthesize)
)
),
g.edge_from(synthesize).to(g.end_node),
)
return g.build()

19
haiku_rag_slim/haiku/rag/research/stream.py
View file

@ -124,7 +124,6 @@ class ResearchStream:
async def stream_research_graph(
graph,
start,
state: "ResearchState",
deps,
) -> AsyncIterator[ResearchStreamEvent]:
@ -132,7 +131,7 @@ async def stream_research_graph(
from contextlib import suppress
from haiku.rag.research.state import ResearchDeps # Local import to avoid cycle
from haiku.rag.research.state import ResearchDeps
if not isinstance(deps, ResearchDeps):
raise TypeError("deps must be an instance of ResearchDeps")
@ -142,25 +141,13 @@ async def stream_research_graph(
async def _execute() -> None:
try:
report = None
try:
result = await graph.run(start, state=state, deps=deps)
report = result.output
except Exception:
from pydantic_graph import End
async with graph.iter(start, state=state, deps=deps) as run:
node = run.next_node
while not isinstance(node, End):
node = await run.next(node)
if run.result:
report = run.result.output
report = await graph.run(state=state, deps=deps)
if report is None:
raise RuntimeError("Graph did not produce a report")
stream.report(report, state)
except Exception as exc: # pragma: no cover - defensive path
except Exception as exc:
stream.error(exc, state)
finally:
await stream.close()

51
tests/test_deep_qa.py
View file

@ -5,14 +5,21 @@ 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.nodes import DeepQAPlanNode
from haiku.rag.qa.deep.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."""
graph = build_deep_qa_graph()
# Mock get_model to return TestModel which generates valid schema-compliant data
def test_model_factory(provider, model):
return TestModel()
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.qa.deep.graph.get_model", test_model_factory)
graph = build_deep_qa_graph(provider="test", model="test")
state = DeepQAState(
context=DeepQAContext(
@ -25,20 +32,12 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
client = HaikuRAG(temp_db_path)
deps = DeepQADeps(client=client, console=None)
# Mock get_model to return TestModel which generates valid schema-compliant data
def test_model_factory(provider, model):
return TestModel()
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.qa.deep.nodes.get_model", test_model_factory)
start = DeepQAPlanNode(provider="test", model="test")
result = await graph.run(start_node=start, state=state, deps=deps)
result = await graph.run(state=state, deps=deps)
# TestModel will generate valid structured output based on schemas
assert result.output.answer is not None
assert isinstance(result.output.answer, str)
assert isinstance(result.output.sources, list)
assert result.answer is not None
assert isinstance(result.answer, str)
assert isinstance(result.sources, list)
client.close()
@ -46,7 +45,15 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
@pytest.mark.asyncio
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()
# Mock get_model to return TestModel
def test_model_factory(provider, model):
return TestModel()
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.qa.deep.graph.get_model", test_model_factory)
graph = build_deep_qa_graph(provider="test", model="test")
state = DeepQAState(
context=DeepQAContext(original_question="What is Python?", use_citations=True),
@ -57,20 +64,12 @@ async def test_deep_qa_with_citations(monkeypatch, temp_db_path):
client = HaikuRAG(temp_db_path)
deps = DeepQADeps(client=client, console=None)
# Mock get_model to return TestModel
def test_model_factory(provider, model):
return TestModel()
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.qa.deep.nodes.get_model", test_model_factory)
start = DeepQAPlanNode(provider="test", model="test")
result = await graph.run(start_node=start, state=state, deps=deps)
result = await graph.run(state=state, deps=deps)
# 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)
assert result.answer is not None
assert isinstance(result.sources, list)
client.close()

5
tests/test_research_graph.py
View file

@ -1,11 +1,12 @@
import asyncio
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import ResearchState, build_research_graph
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.state import ResearchState
def test_build_graph_and_state():
graph = build_research_graph()
graph = build_research_graph(provider="openai", model="gpt-4")
assert graph is not None
state = ResearchState(

35
tests/test_research_graph_integration.py
View file

@ -2,21 +2,25 @@ import pytest
from pydantic_ai.models.test import TestModel
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.nodes.plan import PlanNode
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import (
ResearchDeps,
ResearchState,
build_research_graph,
)
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.models import ResearchReport
from haiku.rag.research.state import ResearchDeps, ResearchState
from haiku.rag.research.stream import stream_research_graph
@pytest.mark.asyncio
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()
# Mock get_model to return TestModel which generates valid schema-compliant data
def test_model_factory(provider, model):
return TestModel()
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.research.graph.get_model", test_model_factory)
graph = build_research_graph(provider="test", model="test")
state = ResearchState(
context=ResearchContext(original_question="What is haiku.rag?"),
@ -29,24 +33,9 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
client = HaikuRAG(temp_db_path)
deps = ResearchDeps(client=client, console=None)
# Mock get_model to return TestModel which generates valid schema-compliant data
# Need to patch in all modules that import it
def test_model_factory(provider, model):
return TestModel()
monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.nodes.plan.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.nodes.search.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.nodes.analysis.get_model", test_model_factory)
monkeypatch.setattr(
"haiku.rag.graph.nodes.synthesize.get_model", test_model_factory
)
start = PlanNode(provider="test", model="test")
collected = []
report = None
async for event in stream_research_graph(graph, start, state, deps):
async for event in stream_research_graph(graph, state, deps):
collected.append(event)
if event.type == "report":
report = event.report