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

Research agent as a pydantic graph

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-19 13:40:52 +03:00
parent 5faa201208
commit cba8ad7696
No known key found for this signature in database
28 changed files with 709 additions and 885 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

71
docs/agents.md
View file

@ -36,50 +36,57 @@ answer = await agent.answer("What is climate change?")
print(answer)
```
### Research Multi‑Agent
### Research Graph
The research workflow coordinates specialized agents to plan, search, analyze, and synthesize a comprehensive answer. It is designed for deeper questions that benefit from iterative investigation and structured reporting.
The research workflow is now implemented as a typed pydantic‑graph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report — with clear stop conditions and shared state.
Components:
Key nodes:
- Orchestrator: Plans, coordinates, and loops until confidence is sufficient
- Presearch Survey: Runs a quick KB scan and summarizes relevant chunk text to
ground the initial plan (plain-text summary; no URIs or scores)
- Search Specialist: Performs targeted RAG searches and answers sub‑questions
- Analysis & Evaluation: Extracts insights, identifies gaps, proposes new questions
- Synthesis: Produces a final structured research report
- Plan: builds up to 3 standalone sub‑questions (uses an internal presearch tool)
- Search (batched): answers sub‑questions using the KB with minimal, verbatim context
- Evaluate: extracts insights, proposes new questions, and checks sufficiency/confidence
- Synthesize: generates a final structured report
Primary models:
- `ResearchPlan` — produced by the orchestrator when planning
- `main_question: str`
- `sub_questions: list[str]` (standalone, self‑contained queries)
- `SearchAnswer` — produced by the search specialist for each sub‑question
- `query: str` — the executed sub‑question
- `answer: str` — the agent’s answer grounded in retrieved context
- `context: list[str]` — minimal verbatim snippets used for the answer
- `sources: list[str]` — document URIs aligned with `context`
- `EvaluationResult` — insights, new standalone questions, sufficiency & confidence
- `ResearchReport` — the final synthesized report
- `SearchAnswer` — one per sub‑question (query, answer, context, sources)
- `EvaluationResult` — insights, new questions, sufficiency, confidence
- `ResearchReport` — final report (title, executive summary, findings, conclusions, …)
CLI usage:
```bash
haiku-rag research "How does haiku.rag organize and query documents?" \
--max-iterations 2 \
--confidence-threshold 0.8 \
--max-concurrency 3 \
--verbose
```
Python usage:
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.research import ResearchOrchestrator
client = HaikuRAG(path_to_db)
orchestrator = ResearchOrchestrator(provider="ollama", model="gpt-oss")
report = await orchestrator.conduct_research(
question="What are the main drivers and recent trends of global temperature anomalies since 1990?",
client=client,
max_iterations=2,
confidence_threshold=0.8,
verbose=True,
from haiku.rag.research import (
ResearchContext,
ResearchDeps,
ResearchState,
build_research_graph,
PlanNode,
)
print(report.title)
print(report.executive_summary)
async with HaikuRAG(path_to_db) as client:
graph = build_research_graph()
state = ResearchState(
question="What are the main drivers and trends of global temperature anomalies since 1990?",
context=ResearchContext(original_question=... ),
max_iterations=2,
confidence_threshold=0.8,
max_concurrency=3,
)
deps = ResearchDeps(client=client)
result = await graph.run(PlanNode(provider=None, model=None), state=state, deps=deps)
report = result.output
print(report.title)
print(report.executive_summary)
```

18
docs/cli.md
View file

@ -84,6 +84,24 @@ haiku-rag ask "Who is the author of haiku.rag?" --cite
The QA agent will search your documents for relevant information and provide a comprehensive answer. With `--cite`, responses include citations showing which documents were used.
## Research
Run the multi-step research graph:
```bash
haiku-rag research "How does haiku.rag organize and query documents?" \
--max-iterations 2 \
--confidence-threshold 0.8 \
--max-concurrency 3 \
--verbose
```
Flags:
- `--max-iterations, -n`: maximum search/evaluate cycles (default: 3)
- `--confidence-threshold`: stop once evaluation confidence meets/exceeds this (default: 0.8)
- `--max-concurrency`: number of sub-questions searched in parallel each iteration (default: 3)
- `--verbose`: show planning, searching previews, evaluation summary, and stop reason
## Server
Start the MCP server:

2
pyproject.toml
View file

@ -29,6 +29,7 @@ dependencies = [
"lancedb>=0.25.0",
"pydantic>=2.11.9",
"pydantic-ai>=1.0.8",
"pydantic-graph>=1.0.8",
"python-dotenv>=1.1.1",
"rich>=14.1.0",
"tiktoken>=0.11.0",
@ -90,6 +91,7 @@ line-ending = "auto"
[tool.pyright]
venvPath = "."
venv = ".venv"
pythonVersion = "3.12"
[tool.pytest.ini_options]
asyncio_default_fixture_loop_scope = "session"

57
src/haiku/rag/app.py
View file

@ -9,7 +9,13 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.mcp import create_mcp_server
from haiku.rag.monitor import FileWatcher
from haiku.rag.research.orchestrator import ResearchOrchestrator
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import (
PlanNode,
ResearchDeps,
ResearchState,
build_research_graph,
)
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
@ -80,28 +86,53 @@ class HaikuRAGApp:
self.console.print(f"[red]Error: {e}[/red]")
async def research(
self, question: str, max_iterations: int = 3, verbose: bool = False
self,
question: str,
max_iterations: int = 3,
confidence_threshold: float = 0.8,
max_concurrency: int = 3,
verbose: bool = False,
):
"""Run multi-agent research on a question."""
"""Run research via the pydantic-graph pipeline (default)."""
async with HaikuRAG(db_path=self.db_path) as client:
try:
# Create orchestrator with default config or fallback to QA
orchestrator = ResearchOrchestrator()
if verbose:
self.console.print(
f"[bold cyan]Starting research with {orchestrator.provider}:{orchestrator.model}[/bold cyan]"
)
self.console.print("[bold cyan]Starting research[/bold cyan]")
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
# Conduct research
report = await orchestrator.conduct_research(
graph = build_research_graph()
state = ResearchState(
question=question,
client=client,
context=ResearchContext(original_question=question),
max_iterations=max_iterations,
verbose=verbose,
confidence_threshold=confidence_threshold,
max_concurrency=max_concurrency,
)
deps = ResearchDeps(
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,
)
# Prefer graph.run; fall back to iter if unavailable
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
if report is None:
raise RuntimeError("Graph did not produce a report")
# Display the report
self.console.print("[bold green]Research Report[/bold green]")

12
src/haiku/rag/cli.py
View file

@ -250,6 +250,16 @@ def research(
"-n",
help="Maximum search/analyze iterations",
),
confidence_threshold: float = typer.Option(
0.8,
"--confidence-threshold",
help="Minimum confidence (0-1) to stop",
),
max_concurrency: int = typer.Option(
3,
"--max-concurrency",
help="Max concurrent searches per iteration (planned)",
),
db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db",
@ -266,6 +276,8 @@ def research(
app.research(
question=question,
max_iterations=max_iterations,
confidence_threshold=confidence_threshold,
max_concurrency=max_concurrency,
verbose=verbose,
)
)

2
src/haiku/rag/reranking/mxbai.py
View file

@ -1,4 +1,4 @@
from mxbai_rerank import MxbaiRerankV2
from mxbai_rerank import MxbaiRerankV2 # pyright: ignore[reportMissingImports]
from haiku.rag.config import Config
from haiku.rag.reranking.base import RerankerBase

37
src/haiku/rag/research/__init__.py
View file

@ -1,37 +1,20 @@
"""Multi-agent research workflow for advanced RAG queries."""
from haiku.rag.research.base import (
BaseResearchAgent,
ResearchOutput,
SearchAnswer,
SearchResult,
)
from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.research.evaluation_agent import (
AnalysisEvaluationAgent,
EvaluationResult,
from haiku.rag.research.graph import (
PlanNode,
ResearchDeps,
ResearchState,
build_research_graph,
)
from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan
from haiku.rag.research.presearch_agent import PresearchSurveyAgent
from haiku.rag.research.search_agent import SearchSpecialistAgent
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
__all__ = [
# Base classes
"BaseResearchAgent",
"ResearchDependencies",
"ResearchContext",
"SearchResult",
"ResearchOutput",
# Specialized agents
"SearchAnswer",
"SearchSpecialistAgent",
"PresearchSurveyAgent",
"AnalysisEvaluationAgent",
"EvaluationResult",
"SynthesisAgent",
"ResearchReport",
# Orchestrator
"ResearchOrchestrator",
"ResearchPlan",
"ResearchDeps",
"ResearchState",
"PlanNode",
"build_research_graph",
]

130
src/haiku/rag/research/base.py
View file

@ -1,130 +0,0 @@
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.output import ToolOutput
from pydantic_ai.providers.ollama import OllamaProvider
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.run import AgentRunResult
from haiku.rag.config import Config
if TYPE_CHECKING:
from haiku.rag.research.dependencies import ResearchDependencies
class BaseResearchAgent[T](ABC):
"""Base class for all research agents."""
def __init__(
self,
provider: str,
model: str,
output_type: type[T],
):
self.provider = provider
self.model = model
self.output_type = output_type
model_obj = self._get_model(provider, model)
# Import deps type lazily to avoid circular import during module load
from haiku.rag.research.dependencies import ResearchDependencies
# If the agent is expected to return plain text, pass `str` directly.
# Otherwise, wrap the model with ToolOutput for robust tool-handling retries.
agent_output_type: Any
if self.output_type is str: # plain text output
agent_output_type = str
else:
agent_output_type = ToolOutput(self.output_type, max_retries=3)
self._agent = Agent(
model=model_obj,
deps_type=ResearchDependencies,
output_type=agent_output_type,
instructions=self.get_system_prompt(),
retries=3,
)
# Register tools
self.register_tools()
def _get_model(self, provider: str, model: str):
"""Get the appropriate model object for the provider."""
if provider == "ollama":
return OpenAIChatModel(
model_name=model,
provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
)
elif provider == "vllm":
return OpenAIChatModel(
model_name=model,
provider=OpenAIProvider(
base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
api_key="none",
),
)
else:
# For all other providers, use the provider:model format
return f"{provider}:{model}"
@abstractmethod
def get_system_prompt(self) -> str:
"""Return the system prompt for this agent."""
pass
def register_tools(self) -> None:
"""Register agent-specific tools."""
pass
async def run(
self, prompt: str, deps: "ResearchDependencies", **kwargs
) -> AgentRunResult[T]:
"""Execute the agent."""
return await self._agent.run(prompt, deps=deps, **kwargs)
@property
def agent(self) -> Agent[Any, T]:
"""Access the underlying Pydantic AI agent."""
return self._agent
class SearchResult(BaseModel):
"""Standard search result format."""
content: str
score: float
document_uri: str
metadata: dict[str, Any] = Field(default_factory=dict)
class ResearchOutput(BaseModel):
"""Standard research output format."""
summary: str
detailed_findings: list[str]
sources: list[str]
confidence: float
class SearchAnswer(BaseModel):
"""Structured output for the SearchSpecialist agent."""
query: str = Field(description="The search query that was performed")
answer: str = Field(description="The answer generated based on the context")
context: list[str] = Field(
description=(
"Only the minimal set of relevant snippets (verbatim) that directly "
"support the answer"
)
)
sources: list[str] = Field(
description=(
"Document URIs corresponding to the snippets actually used in the"
" answer (one URI per snippet; omit if none)"
),
default_factory=list,
)

53
src/haiku/rag/research/common.py Normal file
View file

@ -0,0 +1,53 @@
from typing import Any
from pydantic_ai import format_as_xml
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.ollama import OllamaProvider
from pydantic_ai.providers.openai import OpenAIProvider
from haiku.rag.config import Config
from haiku.rag.research.dependencies import ResearchContext
def get_model(provider: str, model: str) -> Any:
if provider == "ollama":
return OpenAIChatModel(
model_name=model,
provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
)
elif provider == "vllm":
return OpenAIChatModel(
model_name=model,
provider=OpenAIProvider(
base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
api_key="none",
),
)
else:
return f"{provider}:{model}"
def log(console, msg: str) -> None:
if console:
console.print(msg)
def format_context_for_prompt(context: ResearchContext) -> str:
"""Format the research context as XML for inclusion in prompts."""
context_data = {
"original_question": context.original_question,
"unanswered_questions": context.sub_questions,
"qa_responses": [
{
"question": qa.query,
"answer": qa.answer,
"context_snippets": qa.context,
"sources": qa.sources, # pyright: ignore[reportAttributeAccessIssue]
}
for qa in context.qa_responses
],
"insights": context.insights,
"gaps": context.gaps,
}
return format_as_xml(context_data, root_tag="research_context")

28
src/haiku/rag/research/dependencies.py
View file

@ -1,9 +1,8 @@
from pydantic import BaseModel, Field
from pydantic_ai import format_as_xml
from rich.console import Console
from haiku.rag.client import HaikuRAG
from haiku.rag.research.base import SearchAnswer
from haiku.rag.research.models import SearchAnswer
class ResearchContext(BaseModel):
@ -13,7 +12,7 @@ class ResearchContext(BaseModel):
sub_questions: list[str] = Field(
default_factory=list, description="Decomposed sub-questions"
)
qa_responses: list["SearchAnswer"] = Field(
qa_responses: list[SearchAnswer] = Field(
default_factory=list, description="Structured QA pairs used during research"
)
insights: list[str] = Field(
@ -23,7 +22,7 @@ class ResearchContext(BaseModel):
default_factory=list, description="Identified information gaps"
)
def add_qa_response(self, qa: "SearchAnswer") -> None:
def add_qa_response(self, qa: SearchAnswer) -> None:
"""Add a structured QA response (minimal context already included)."""
self.qa_responses.append(qa)
@ -46,24 +45,3 @@ class ResearchDependencies(BaseModel):
client: HaikuRAG = Field(description="RAG client for document operations")
context: ResearchContext = Field(description="Shared research context")
console: Console | None = None
def _format_context_for_prompt(context: ResearchContext) -> str:
"""Format the research context as XML for inclusion in prompts."""
context_data = {
"original_question": context.original_question,
"unanswered_questions": context.sub_questions,
"qa_responses": [
{
"question": qa.query,
"answer": qa.answer,
"context_snippets": qa.context,
"sources": qa.sources,
}
for qa in context.qa_responses
],
"insights": context.insights,
"gaps": context.gaps,
}
return format_as_xml(context_data, root_tag="research_context")

85
src/haiku/rag/research/evaluation_agent.py
View file

@ -1,85 +0,0 @@
from pydantic import BaseModel, Field
from pydantic_ai.run import AgentRunResult
from haiku.rag.research.base import BaseResearchAgent
from haiku.rag.research.dependencies import (
ResearchDependencies,
_format_context_for_prompt,
)
from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
class EvaluationResult(BaseModel):
"""Result of analysis and evaluation."""
key_insights: list[str] = Field(
description="Main insights extracted from the research so far"
)
new_questions: list[str] = Field(
description="New sub-questions to add to the research (max 3)",
max_length=3,
default=[],
)
confidence_score: float = Field(
description="Confidence level in the completeness of research (0-1)",
ge=0.0,
le=1.0,
)
is_sufficient: bool = Field(
description="Whether the research is sufficient to answer the original question"
)
reasoning: str = Field(
description="Explanation of why the research is or isn't complete"
)
class AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]):
"""Agent that analyzes findings and evaluates research completeness."""
def __init__(self, provider: str, model: str) -> None:
super().__init__(provider, model, output_type=EvaluationResult)
async def run(
self, prompt: str, deps: ResearchDependencies, **kwargs
) -> AgentRunResult[EvaluationResult]:
console = deps.console
if console:
console.print(
"\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]"
)
# Format context for the evaluation agent
context_xml = _format_context_for_prompt(deps.context)
evaluation_prompt = f"""Analyze all gathered information and evaluate the completeness of research.
{context_xml}
Evaluate the research progress for the original question and identify any remaining gaps."""
result = await super().run(evaluation_prompt, deps, **kwargs)
output = result.output
# Store insights
for insight in output.key_insights:
deps.context.add_insight(insight)
# Add new questions to the sub-questions list
for new_q in output.new_questions:
if new_q not in deps.context.sub_questions:
deps.context.sub_questions.append(new_q)
if console:
if output.key_insights:
console.print(" [bold]Key insights:[/bold]")
for insight in output.key_insights:
console.print(f" • {insight}")
console.print(
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]"
)
status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
console.print(f" Sufficient: {status}")
return result
def get_system_prompt(self) -> str:
return EVALUATION_AGENT_PROMPT

29
src/haiku/rag/research/graph.py Normal file
View file

@ -0,0 +1,29 @@
from pydantic_graph import Graph
from haiku.rag.research.models import ResearchReport
from haiku.rag.research.nodes.evaluate import EvaluateNode
from haiku.rag.research.nodes.plan import PlanNode
from haiku.rag.research.nodes.search import SearchDispatchNode
from haiku.rag.research.nodes.synthesize import SynthesizeNode
from haiku.rag.research.state import ResearchDeps, ResearchState
__all__ = [
"PlanNode",
"SearchDispatchNode",
"EvaluateNode",
"SynthesizeNode",
"ResearchState",
"ResearchDeps",
"build_research_graph",
]
def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]:
return Graph(
nodes=[
PlanNode,
SearchDispatchNode,
EvaluateNode,
SynthesizeNode,
]
)

70
src/haiku/rag/research/models.py Normal file
View file

@ -0,0 +1,70 @@
from pydantic import BaseModel, Field
class ResearchPlan(BaseModel):
main_question: str
sub_questions: list[str]
class SearchAnswer(BaseModel):
"""Structured output for the SearchSpecialist agent."""
query: str = Field(description="The search query that was performed")
answer: str = Field(description="The answer generated based on the context")
context: list[str] = Field(
description=(
"Only the minimal set of relevant snippets (verbatim) that directly "
"support the answer"
)
)
sources: list[str] = Field(
description=(
"Document URIs corresponding to the snippets actually used in the"
" answer (one URI per snippet; omit if none)"
),
default_factory=list,
)
class EvaluationResult(BaseModel):
"""Result of analysis and evaluation."""
key_insights: list[str] = Field(
description="Main insights extracted from the research so far"
)
new_questions: list[str] = Field(
description="New sub-questions to add to the research (max 3)",
max_length=3,
default=[],
)
confidence_score: float = Field(
description="Confidence level in the completeness of research (0-1)",
ge=0.0,
le=1.0,
)
is_sufficient: bool = Field(
description="Whether the research is sufficient to answer the original question"
)
reasoning: str = Field(
description="Explanation of why the research is or isn't complete"
)
class ResearchReport(BaseModel):
"""Final research report structure."""
title: str = Field(description="Concise title for the research")
executive_summary: str = Field(description="Brief overview of key findings")
main_findings: list[str] = Field(
description="Primary research findings with supporting evidence"
)
conclusions: list[str] = Field(description="Evidence-based conclusions")
limitations: list[str] = Field(
description="Limitations of the current research", default=[]
)
recommendations: list[str] = Field(
description="Actionable recommendations based on findings", default=[]
)
sources_summary: str = Field(
description="Summary of sources used and their reliability"
)

81
src/haiku/rag/research/nodes/evaluate.py Normal file
View file

@ -0,0 +1,81 @@
from dataclasses import dataclass
from pydantic_ai import Agent
from pydantic_graph import BaseNode, GraphRunContext
from haiku.rag.research.common import format_context_for_prompt, get_model, log
from haiku.rag.research.dependencies import (
ResearchDependencies,
)
from haiku.rag.research.models import EvaluationResult, ResearchReport
from haiku.rag.research.nodes.synthesize import SynthesizeNode
from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
from haiku.rag.research.state import ResearchDeps, ResearchState
@dataclass
class EvaluateNode(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.console,
"\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]",
)
agent = Agent(
model=get_model(self.provider, self.model),
output_type=EvaluationResult,
instructions=EVALUATION_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Analyze gathered information and evaluate completeness for the original question.\n\n"
f"{context_xml}"
)
agent_deps = ResearchDependencies(
client=deps.client, context=state.context, console=deps.console
)
eval_result = await agent.run(prompt, deps=agent_deps)
output = eval_result.output
for insight in output.key_insights:
state.context.add_insight(insight)
for new_q in output.new_questions:
if new_q not in state.sub_questions:
state.sub_questions.append(new_q)
state.last_eval = output
state.iterations += 1
if deps.console:
if output.key_insights:
deps.console.print(" [bold]Key insights:[/bold]")
for ins in output.key_insights:
deps.console.print(f" • {ins}")
deps.console.print(
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]"
)
status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
deps.console.print(f" Sufficient: {status}")
from haiku.rag.research.nodes.search import SearchDispatchNode
if (
output.is_sufficient
and output.confidence_score >= state.confidence_threshold
) or state.iterations >= state.max_iterations:
if deps.console:
deps.console.print("\n[bold green]✅ Stopping research.[/bold green]")
return SynthesizeNode(self.provider, self.model)
return SearchDispatchNode(self.provider, self.model)

64
src/haiku/rag/research/nodes/plan.py Normal file
View file

@ -0,0 +1,64 @@
from dataclasses import dataclass
from pydantic_ai import Agent, RunContext
from pydantic_graph import BaseNode, GraphRunContext
from haiku.rag.research.common import get_model, log
from haiku.rag.research.dependencies import ResearchDependencies
from haiku.rag.research.models import ResearchPlan, ResearchReport
from haiku.rag.research.nodes.search import SearchDispatchNode
from haiku.rag.research.prompts import ORCHESTRATOR_PROMPT
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.console, "\n[bold cyan]📋 Creating research plan...[/bold cyan]")
plan_agent = Agent(
model=get_model(self.provider, self.model),
output_type=ResearchPlan,
instructions=(
ORCHESTRATOR_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.question}"
)
agent_deps = ResearchDependencies(
client=deps.client, context=state.context, console=deps.console
)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.sub_questions = list(plan_result.output.sub_questions)
if deps.console:
deps.console.print("\n[bold green]✅ Research Plan Created:[/bold green]")
deps.console.print(f" [bold]Main Question:[/bold] {state.question}")
deps.console.print(" [bold]Sub-questions:[/bold]")
for i, sq in enumerate(state.sub_questions, 1):
deps.console.print(f" {i}. {sq}")
return SearchDispatchNode(self.provider, self.model)

92
src/haiku/rag/research/nodes/search.py Normal file
View file

@ -0,0 +1,92 @@
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.research.common import get_model, log
from haiku.rag.research.dependencies import ResearchDependencies
from haiku.rag.research.models import ResearchReport, SearchAnswer
from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT
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.sub_questions:
from haiku.rag.research.nodes.evaluate import EvaluateNode
return EvaluateNode(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.sub_questions and len(batch) < take:
batch.append(state.sub_questions.pop(0))
async def answer_one(sub_q: str) -> SearchAnswer | None:
if deps.console:
deps.console.print(
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_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
)
try:
result = await agent.run(sub_q, deps=agent_deps)
except Exception as e:
log(deps.console, 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)
if deps.console:
preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
deps.console.log(f" [green]✓[/green] {preview}")
return SearchDispatchNode(self.provider, self.model)

51
src/haiku/rag/research/nodes/synthesize.py Normal file
View file

@ -0,0 +1,51 @@
from dataclasses import dataclass
from pydantic_ai import Agent
from pydantic_graph import BaseNode, End, GraphRunContext
from haiku.rag.research.common import format_context_for_prompt, get_model, log
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.console,
"\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
)
result = await agent.run(prompt, deps=agent_deps)
log(deps.console, "[bold green]✅ Research complete![/bold green]")
return End(result.output)

170
src/haiku/rag/research/orchestrator.py
View file

@ -1,170 +0,0 @@
from typing import Any
from pydantic import BaseModel, Field
from pydantic_ai.run import AgentRunResult
from rich.console import Console
from haiku.rag.config import Config
from haiku.rag.research.base import BaseResearchAgent
from haiku.rag.research.dependencies import (
ResearchContext,
ResearchDependencies,
)
from haiku.rag.research.evaluation_agent import (
AnalysisEvaluationAgent,
EvaluationResult,
)
from haiku.rag.research.presearch_agent import PresearchSurveyAgent
from haiku.rag.research.prompts import ORCHESTRATOR_PROMPT
from haiku.rag.research.search_agent import SearchSpecialistAgent
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
class ResearchPlan(BaseModel):
"""Research execution plan."""
main_question: str = Field(description="The main research question")
sub_questions: list[str] = Field(
description="Decomposed sub-questions to investigate (max 3)", max_length=3
)
class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
"""Orchestrator agent that coordinates the research workflow."""
def __init__(
self,
provider: str | None = Config.RESEARCH_PROVIDER,
model: str | None = None,
):
# Use provided values or fall back to config defaults
provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER
model = model or Config.RESEARCH_MODEL or Config.QA_MODEL
super().__init__(provider, model, output_type=ResearchPlan)
self.search_agent: SearchSpecialistAgent = SearchSpecialistAgent(
provider, model
)
self.presearch_agent: PresearchSurveyAgent = PresearchSurveyAgent(
provider, model
)
self.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent(
provider, model
)
self.synthesis_agent: SynthesisAgent = SynthesisAgent(provider, model)
def get_system_prompt(self) -> str:
return ORCHESTRATOR_PROMPT
def _should_stop_research(
self,
evaluation_result: AgentRunResult[EvaluationResult],
confidence_threshold: float,
) -> bool:
"""Determine if research should stop based on evaluation."""
result = evaluation_result.output
return result.is_sufficient and result.confidence_score >= confidence_threshold
async def conduct_research(
self,
question: str,
client: Any,
max_iterations: int = 3,
confidence_threshold: float = 0.8,
verbose: bool = False,
) -> ResearchReport:
"""Conduct comprehensive research on a question.
Args:
question: The research question to investigate
client: HaikuRAG client for document operations
max_iterations: Maximum number of search-analyze-clarify cycles
confidence_threshold: Minimum confidence level to stop research (0-1)
verbose: If True, print progress and intermediate results
Returns:
ResearchReport with comprehensive findings
"""
# Initialize context
context = ResearchContext(original_question=question)
deps = ResearchDependencies(client=client, context=context)
if verbose:
deps.console = Console()
console = deps.console
# Create initial research plan
if console:
console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
# Run a simple presearch survey to summarize KB context
presearch_result = await self.presearch_agent.run(question, deps=deps)
plan_prompt = (
"Create a research plan for the main question below.\n\n"
f"Main question: {question}\n\n"
"Use this brief presearch summary to inform the plan. Focus the 3 sub-questions "
"on the most important aspects not already obvious from the current KB context.\n\n"
f"{presearch_result.output}"
)
plan_result: AgentRunResult[ResearchPlan] = await self.run(
plan_prompt, deps=deps
)
context.sub_questions = plan_result.output.sub_questions
if console:
console.print("\n[bold green]✅ Research Plan Created:[/bold green]")
console.print(
f" [bold]Main Question:[/bold] {plan_result.output.main_question}"
)
console.print(" [bold]Sub-questions:[/bold]")
for i, sq in enumerate(plan_result.output.sub_questions, 1):
console.print(f" {i}. {sq}")
# Execute research iterations
for iteration in range(max_iterations):
if console:
console.rule(
f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]"
)
# Check if we have questions to search
if not context.sub_questions:
if console:
console.print(
"[yellow]No more questions to explore. Concluding research.[/yellow]"
)
break
# Use current sub-questions for this iteration
questions_to_search = context.sub_questions[:]
# Search phase - answer all questions in this iteration
if console:
console.print(
f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]"
)
for search_question in questions_to_search:
await self.search_agent.run(search_question, deps=deps)
# Analysis and Evaluation phase
evaluation_result = await self.evaluation_agent.run("", deps=deps)
# Check if research is sufficient
if self._should_stop_research(evaluation_result, confidence_threshold):
if console:
console.print(
f"\n[bold green]✅ Stopping research:[/bold green] {evaluation_result.output.reasoning}"
)
break
# Generate final report
report_result: AgentRunResult[ResearchReport] = await self.synthesis_agent.run(
"", deps=deps
)
return report_result.output

39
src/haiku/rag/research/presearch_agent.py
View file

@ -1,39 +0,0 @@
from pydantic_ai import RunContext
from pydantic_ai.run import AgentRunResult
from haiku.rag.research.base import BaseResearchAgent
from haiku.rag.research.dependencies import ResearchDependencies
from haiku.rag.research.prompts import PRESEARCH_AGENT_PROMPT
class PresearchSurveyAgent(BaseResearchAgent[str]):
"""Presearch agent that gathers verbatim context and summarizes it."""
def __init__(self, provider: str, model: str) -> None:
super().__init__(provider, model, str)
async def run(
self, prompt: str, deps: ResearchDependencies, **kwargs
) -> AgentRunResult[str]:
console = deps.console
if console:
console.print(
"\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]"
)
return await super().run(prompt, deps, **kwargs)
def get_system_prompt(self) -> str:
return PRESEARCH_AGENT_PROMPT
def register_tools(self) -> None:
@self.agent.tool
async def gather_context(
ctx: RunContext[ResearchDependencies],
query: str,
limit: int = 6,
) -> str:
"""Return verbatim concatenation of relevant chunk texts."""
results = await ctx.deps.client.search(query, limit=limit)
expanded = await ctx.deps.client.expand_context(results)
return "\n\n".join(chunk.content for chunk, _ in expanded)

69
src/haiku/rag/research/search_agent.py
View file

@ -1,69 +0,0 @@
from pydantic_ai import RunContext
from pydantic_ai.format_prompt import format_as_xml
from pydantic_ai.run import AgentRunResult
from haiku.rag.research.base import BaseResearchAgent, SearchAnswer
from haiku.rag.research.dependencies import ResearchDependencies
from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT
class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
"""Agent specialized in answering questions using RAG search."""
def __init__(self, provider: str, model: str) -> None:
super().__init__(provider, model, output_type=SearchAnswer)
async def run(
self, prompt: str, deps: ResearchDependencies, **kwargs
) -> AgentRunResult[SearchAnswer]:
"""Execute the agent and persist the QA pair in shared context.
Pydantic AI enforces `SearchAnswer` as the output model; we just store
the QA response with the last search results as sources.
"""
console = deps.console
if console:
console.print(f"\t{prompt}")
result = await super().run(prompt, deps, **kwargs)
deps.context.add_qa_response(result.output)
deps.context.sub_questions.remove(prompt)
if console:
answer = result.output.answer
answer_preview = answer[:150] + "…" if len(answer) > 150 else answer
console.log(f"\n [green]✓[/green] {answer_preview}")
return result
def get_system_prompt(self) -> str:
return SEARCH_AGENT_PROMPT
def register_tools(self) -> None:
"""Register search-specific tools."""
@self.agent.tool
async def search_and_answer(
ctx: RunContext[ResearchDependencies],
query: str,
limit: int = 5,
) -> str:
"""Search the KB and return a concise context pack."""
search_results = await ctx.deps.client.search(query, limit=limit)
expanded = await ctx.deps.client.expand_context(search_results)
snippet_entries = [
{
"text": chunk.content,
"score": score,
"document_uri": (chunk.document_uri or ""),
}
for chunk, score in expanded
]
# Return an XML-formatted payload with the question and snippets.
if snippet_entries:
return format_as_xml(snippet_entries, root_tag="snippets")
else:
return (
f"No relevant information found in the knowledge base for: {query}"
)

25
src/haiku/rag/research/state.py Normal file
View file

@ -0,0 +1,25 @@
from dataclasses import dataclass, field
from rich.console import Console
from haiku.rag.client import HaikuRAG
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.models import EvaluationResult
@dataclass
class ResearchDeps:
client: HaikuRAG
console: Console | None = None
@dataclass
class ResearchState:
question: str
context: ResearchContext
sub_questions: list[str] = field(default_factory=list)
iterations: int = 0
max_iterations: int = 3
max_concurrency: int = 3
confidence_threshold: float = 0.8
last_eval: EvaluationResult | None = None

60
src/haiku/rag/research/synthesis_agent.py
View file

@ -1,60 +0,0 @@
from pydantic import BaseModel, Field
from pydantic_ai.run import AgentRunResult
from haiku.rag.research.base import BaseResearchAgent
from haiku.rag.research.dependencies import (
ResearchDependencies,
_format_context_for_prompt,
)
from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT
class ResearchReport(BaseModel):
"""Final research report structure."""
title: str = Field(description="Concise title for the research")
executive_summary: str = Field(description="Brief overview of key findings")
main_findings: list[str] = Field(
description="Primary research findings with supporting evidence"
)
conclusions: list[str] = Field(description="Evidence-based conclusions")
limitations: list[str] = Field(
description="Limitations of the current research", default=[]
)
recommendations: list[str] = Field(
description="Actionable recommendations based on findings", default=[]
)
sources_summary: str = Field(
description="Summary of sources used and their reliability"
)
class SynthesisAgent(BaseResearchAgent[ResearchReport]):
"""Agent specialized in synthesizing research into comprehensive reports."""
def __init__(self, provider: str, model: str) -> None:
super().__init__(provider, model, output_type=ResearchReport)
async def run(
self, prompt: str, deps: ResearchDependencies, **kwargs
) -> AgentRunResult[ResearchReport]:
console = deps.console
if console:
console.print(
"\n[bold cyan]📝 Generating final research report...[/bold cyan]"
)
context_xml = _format_context_for_prompt(deps.context)
synthesis_prompt = f"""Generate a comprehensive research report based on all gathered information.
{context_xml}
Create a detailed report that synthesizes all findings into a coherent response."""
result = await super().run(synthesis_prompt, deps, **kwargs)
if console:
console.print("[bold green]✅ Research complete![/bold green]")
return result
def get_system_prompt(self) -> str:
return SYNTHESIS_AGENT_PROMPT

17
tests/research/test_evaluation_agent.py
View file

@ -1,17 +0,0 @@
from haiku.rag.config import Config
from haiku.rag.research.evaluation_agent import (
AnalysisEvaluationAgent,
EvaluationResult,
)
class TestAnalysisEvaluationAgent:
"""Lean tests for AnalysisEvaluationAgent without LLM mocking."""
def test_agent_initialization(self):
agent = AnalysisEvaluationAgent(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
assert agent.provider == Config.RESEARCH_PROVIDER
assert agent.model == Config.RESEARCH_MODEL
assert agent.output_type == EvaluationResult

189
tests/research/test_orchestrator.py
View file

@ -1,189 +0,0 @@
from unittest.mock import AsyncMock, create_autospec
import pytest
from pydantic_ai.models.test import TestModel
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.research.evaluation_agent import EvaluationResult
from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan
from haiku.rag.research.synthesis_agent import ResearchReport
from haiku.rag.store.models.chunk import Chunk
@pytest.fixture
def test_model():
"""Create a test model for orchestrator testing."""
return TestModel()
@pytest.fixture
def mock_client():
"""Create a mock HaikuRAG client."""
client = create_autospec(HaikuRAG, instance=True)
client.search = AsyncMock()
client.expand_context = AsyncMock()
return client
@pytest.fixture
def research_context():
"""Create a research context."""
return ResearchContext(original_question="What is climate change?")
@pytest.fixture
def research_deps(mock_client, research_context):
"""Create research dependencies."""
return ResearchDependencies(client=mock_client, context=research_context)
def create_mock_chunk(chunk_id: str, content: str, score: float = 0.8):
"""Helper to create mock chunk objects."""
return Chunk(
id=chunk_id,
document_id=f"doc_{chunk_id}",
content=content,
document_uri=f"doc_{chunk_id}.md",
metadata={},
), score
class TestResearchOrchestrator:
"""Test suite for ResearchOrchestrator."""
def test_orchestrator_uses_config_defaults(self):
"""Test that orchestrator uses config defaults when no args provided."""
orchestrator = ResearchOrchestrator()
# Should use RESEARCH_PROVIDER/MODEL if set, else QA_PROVIDER/MODEL
assert orchestrator.provider is not None
assert orchestrator.model is not None
# All agents should use the same provider/model
assert orchestrator.search_agent.provider == orchestrator.provider
assert orchestrator.search_agent.model == orchestrator.model
assert orchestrator.evaluation_agent.provider == orchestrator.provider
assert orchestrator.evaluation_agent.model == orchestrator.model
assert orchestrator.synthesis_agent.provider == orchestrator.provider
assert orchestrator.synthesis_agent.model == orchestrator.model
def test_orchestrator_initialization(self):
"""Test that orchestrator initializes all agents correctly."""
orchestrator = ResearchOrchestrator(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
# Check all agents are initialized
assert orchestrator.search_agent is not None
assert orchestrator.evaluation_agent is not None
assert orchestrator.synthesis_agent is not None
# Check they all use the same provider and model
assert orchestrator.search_agent.provider == Config.RESEARCH_PROVIDER
assert orchestrator.search_agent.model == Config.RESEARCH_MODEL
assert orchestrator.evaluation_agent.provider == Config.RESEARCH_PROVIDER
assert orchestrator.evaluation_agent.model == Config.RESEARCH_MODEL
assert orchestrator.synthesis_agent.provider == Config.RESEARCH_PROVIDER
assert orchestrator.synthesis_agent.model == Config.RESEARCH_MODEL
def test_orchestrator_has_correct_output_type(self):
"""Test that orchestrator's output type is ResearchPlan."""
orchestrator = ResearchOrchestrator(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
assert orchestrator.output_type == ResearchPlan
def test_orchestrator_has_no_tools(self):
"""Test that orchestrator no longer registers tools (direct agent calls now)."""
orchestrator = ResearchOrchestrator(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
# Get the tools from the agent
tools = orchestrator.agent._function_toolset.tools
tool_names = list(tools.keys())
# Should have no tools since we call agents directly now
assert len(tool_names) == 0
def test_should_stop_research_logic(self):
"""Test the stopping logic based on EvaluationResult."""
orchestrator = ResearchOrchestrator(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
# Create mock evaluation results
from unittest.mock import MagicMock
# Sufficient research result
sufficient_result = MagicMock()
sufficient_result.output = EvaluationResult(
key_insights=["Climate is changing", "Human activity is the cause"],
new_questions=[],
confidence_score=0.9,
is_sufficient=True,
reasoning="All aspects covered comprehensively",
)
# Insufficient research result
insufficient_result = MagicMock()
insufficient_result.output = EvaluationResult(
key_insights=["Some data found"],
new_questions=[
"What about economic impacts?",
"Regional variations?",
],
confidence_score=0.4,
is_sufficient=False,
reasoning="Major gaps remain in understanding",
)
# Test with sufficient research (threshold 0.8)
assert orchestrator._should_stop_research(sufficient_result, 0.8)
# Test with insufficient research
assert not orchestrator._should_stop_research(insufficient_result, 0.8)
# Test with high confidence but below threshold
sufficient_result.output.confidence_score = 0.75
assert not orchestrator._should_stop_research(sufficient_result, 0.8)
# Test with is_sufficient=False even with high confidence
insufficient_result.output.confidence_score = 0.95
assert not orchestrator._should_stop_research(insufficient_result, 0.8)
@pytest.mark.asyncio
async def test_conduct_research_workflow(self, test_model, mock_client):
"""Test the basic research workflow using TestModel."""
orchestrator = ResearchOrchestrator(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
# Setup mock client returns
mock_chunks = [
create_mock_chunk("1", "Climate change information"),
]
mock_client.search.return_value = mock_chunks
mock_client.expand_context.return_value = mock_chunks
# Use TestModel for all agents
with orchestrator.agent.override(model=test_model):
with orchestrator.search_agent.agent.override(model=test_model):
with orchestrator.evaluation_agent.agent.override(model=test_model):
with orchestrator.synthesis_agent.agent.override(model=test_model):
# Run the research
report = await orchestrator.conduct_research(
"What is climate change?", mock_client, max_iterations=1
)
# Verify we got a valid report structure
assert isinstance(report, ResearchReport)
assert report.title
assert report.executive_summary
assert isinstance(report.main_findings, list)
assert isinstance(report.conclusions, list)
assert isinstance(report.limitations, list)
assert isinstance(report.recommendations, list)
assert report.sources_summary

14
tests/research/test_search_agent.py
View file

@ -1,14 +0,0 @@
from haiku.rag.config import Config
from haiku.rag.research import SearchAnswer, SearchSpecialistAgent
class TestSearchSpecialistAgent:
"""Lean tests for SearchSpecialistAgent without LLM mocking."""
def test_agent_initialization(self):
agent = SearchSpecialistAgent(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
assert agent.provider == Config.RESEARCH_PROVIDER
assert agent.model == Config.RESEARCH_MODEL
assert agent.output_type is SearchAnswer

14
tests/research/test_synthesis_agent.py
View file

@ -1,14 +0,0 @@
from haiku.rag.config import Config
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
class TestSynthesisAgent:
"""Lean tests for SynthesisAgent without LLM mocking."""
def test_agent_initialization(self):
agent = SynthesisAgent(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
assert agent.provider == Config.RESEARCH_PROVIDER
assert agent.model == Config.RESEARCH_MODEL
assert agent.output_type == ResearchReport

26
tests/test_research_graph.py Normal file
View file

@ -0,0 +1,26 @@
import asyncio
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import ResearchState, build_research_graph
def test_build_graph_and_state():
graph = build_research_graph()
assert graph is not None
state = ResearchState(
question="What are the key features of haiku.rag?",
context=ResearchContext(
original_question="What are the key features of haiku.rag?"
),
max_iterations=1,
confidence_threshold=0.8,
)
assert state.iterations == 0
assert state.sub_questions == []
def test_async_loop_available():
# Ensure an event loop can be created in test env
loop = asyncio.new_event_loop()
loop.close()

89
tests/test_research_graph_integration.py Normal file
View file

@ -0,0 +1,89 @@
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