Merge pull request #64 from ggozad/feat/research-as-graph

Research agents as a graph

.pre-commit-config.yaml

@@ -20,13 +20,3 @@ repos:
     rev: v1.1.399
     hooks:
       - id: pyright
-
-  - repo: https://github.com/RodrigoGonzalez/check-mkdocs
-    rev: v1.2.0
-    hooks:
-      - id: check-mkdocs
-        name: check-mkdocs
-        args: ["--config", "mkdocs.yml"] # Optional, mkdocs.yml is the default
-        # If you have additional plugins or libraries that are not included in
-        # check-mkdocs, add them here
-        additional_dependencies: ["mkdocs-material"]

README.md

@@ -11,6 +11,7 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB.
 - **Local LanceDB**: No external servers required, supports also LanceDB cloud storage, S3, Google Cloud & Azure
 - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI, vLLM
 - **Multiple QA providers**: Any provider/model supported by Pydantic AI
+- **Research graph (multi‑agent)**: Plan → Search → Evaluate → Synthesize with agentic AI
 - **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking
 - **Reranking**: Default search result reranking with MixedBread AI, Cohere, or vLLM
 - **Question answering**: Built-in QA agents on your documents
@@ -38,6 +39,14 @@ haiku-rag ask "Who is the author of haiku.rag?"
 # Ask questions with citations
 haiku-rag ask "Who is the author of haiku.rag?" --cite
 
+# Multi‑agent research (iterative plan/search/evaluate)
+haiku-rag research \
+  "What are the main drivers and trends of global temperature anomalies since 1990?" \
+  --max-iterations 2 \
+  --confidence-threshold 0.8 \
+  --max-concurrency 3 \
+  --verbose
+
 # Rebuild database (re-chunk and re-embed all documents)
 haiku-rag rebuild
 
@@ -53,6 +62,13 @@ haiku-rag serve
 
 ```python
 from haiku.rag.client import HaikuRAG
+from haiku.rag.research import (
+    ResearchContext,
+    ResearchDeps,
+    ResearchState,
+    build_research_graph,
+    PlanNode,
+)
 
 async with HaikuRAG("database.lancedb") as client:
     # Add document
@@ -70,6 +86,25 @@ async with HaikuRAG("database.lancedb") as client:
     # Ask questions with citations
     answer = await client.ask("Who is the author of haiku.rag?", cite=True)
     print(answer)
+
+    # Multi‑agent research pipeline (Plan → Search → Evaluate → Synthesize)
+    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)
+    start = PlanNode(provider=None, model=None)
+    result = await graph.run(start, state=state, deps=deps)
+    report = result.output
+    print(report.title)
+    print(report.executive_summary)
 ```
 
 ## MCP Server

docs/agents.md

@@ -36,50 +36,69 @@ 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 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:
+```mermaid
+---
+title: Research graph
+---
+stateDiagram-v2
+  PlanNode --> SearchDispatchNode
+  SearchDispatchNode --> EvaluateNode
+  EvaluateNode --> SearchDispatchNode
+  EvaluateNode --> SynthesizeNode
+  SynthesizeNode --> [*]
+```
 
-- Orchestrator: Plans, coordinates, and loops until confidence is sufficient
+Key nodes:
-- Presearch Survey: Runs a quick KB scan and summarizes relevant chunk text to
+
-  ground the initial plan (plain-text summary; no URIs or scores)
+- Plan: builds up to 3 standalone sub‑questions (uses an internal presearch tool)
-- Search Specialist: Performs targeted RAG searches and answers sub‑questions
+- Search (batched): answers sub‑questions using the KB with minimal, verbatim context
-- Analysis & Evaluation: Extracts insights, identifies gaps, proposes new questions
+- Evaluate: extracts insights, proposes new questions, and checks sufficiency/confidence
-- Synthesis: Produces a final structured research report
+- Synthesize: generates a final structured report
 
 Primary models:
 
-- `ResearchPlan` — produced by the orchestrator when planning
+- `SearchAnswer` — one per sub‑question (query, answer, context, sources)
-  - `main_question: str`
+- `EvaluationResult` — insights, new questions, sufficiency, confidence
-  - `sub_questions: list[str]` (standalone, self‑contained queries)
+- `ResearchReport` — final report (title, executive summary, findings, conclusions, …)
-- `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
 
+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
+from haiku.rag.research import (
-
+    ResearchContext,
-client = HaikuRAG(path_to_db)
+    ResearchDeps,
-orchestrator = ResearchOrchestrator(provider="ollama", model="gpt-oss")
+    ResearchState,
-
+    build_research_graph,
-report = await orchestrator.conduct_research(
+    PlanNode,
-    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,
 )
 
-print(report.title)
+async with HaikuRAG(path_to_db) as client:
-print(report.executive_summary)
+    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)
 ```

docs/cli.md

@@ -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:

mkdocs.yml

@@ -76,4 +76,8 @@ markdown_extensions:
       use_pygments: true
   - pymdownx.inlinehilite
   - pymdownx.snippets
-  - pymdownx.superfences
+  - pymdownx.superfences:
+      custom_fences:
+        - name: mermaid
+          class: mermaid
+          format: !!python/name:pymdownx.superfences.fence_code_format

pyproject.toml

@@ -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"

src/haiku/rag/app.py

@@ -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 = 1,
+        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(
+                    self.console.print("[bold cyan]Starting research[/bold cyan]")
-                        f"[bold cyan]Starting research with {orchestrator.provider}:{orchestrator.model}[/bold cyan]"
-                    )
                     self.console.print(f"[bold blue]Question:[/bold blue] {question}")
                     self.console.print()
 
-                # Conduct research
+                graph = build_research_graph()
-                report = await orchestrator.conduct_research(
+                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]")
@@ -114,6 +145,12 @@ class HaikuRAGApp:
                 self.console.print(report.executive_summary)
                 self.console.print()
 
+                # Confidence (from last evaluation)
+                if state.last_eval:
+                    conf = state.last_eval.confidence_score  # type: ignore[attr-defined]
+                    self.console.print(f"[bold cyan]Confidence:[/bold cyan] {conf:.1%}")
+                    self.console.print()
+
                 # Main Findings
                 if report.main_findings:
                     self.console.print("[bold cyan]Main Findings:[/bold cyan]")

src/haiku/rag/cli.py

@@ -13,10 +13,10 @@ from haiku.rag.logging import configure_cli_logging
 from haiku.rag.migration import migrate_sqlite_to_lancedb
 from haiku.rag.utils import is_up_to_date
 
-logfire.configure(send_to_logfire="if-token-present")
+if Config.ENV == "development":
-logfire.instrument_pydantic_ai()
+    logfire.configure(send_to_logfire="if-token-present")
-
+    logfire.instrument_pydantic_ai()
-if not Config.ENV == "development":
+else:
     warnings.filterwarnings("ignore")
 
 cli = typer.Typer(
@@ -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(
+        1,
+        "--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,
         )
     )

src/haiku/rag/reranking/mxbai.py

@@ -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

src/haiku/rag/research/__init__.py

@@ -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 (
+from haiku.rag.research.graph import (
-    AnalysisEvaluationAgent,
+    PlanNode,
-    EvaluationResult,
+    ResearchDeps,
+    ResearchState,
+    build_research_graph,
 )
-from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan
+from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
-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
 
 __all__ = [
-    # Base classes
-    "BaseResearchAgent",
     "ResearchDependencies",
     "ResearchContext",
-    "SearchResult",
-    "ResearchOutput",
-    # Specialized agents
     "SearchAnswer",
-    "SearchSpecialistAgent",
-    "PresearchSurveyAgent",
-    "AnalysisEvaluationAgent",
     "EvaluationResult",
-    "SynthesisAgent",
     "ResearchReport",
-    # Orchestrator
+    "ResearchDeps",
-    "ResearchOrchestrator",
+    "ResearchState",
-    "ResearchPlan",
+    "PlanNode",
+    "build_research_graph",
 ]

src/haiku/rag/research/base.py

@@ -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,
-    )

src/haiku/rag/research/common.py

@@ -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")

src/haiku/rag/research/dependencies.py

@@ -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")

src/haiku/rag/research/evaluation_agent.py

@@ -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

src/haiku/rag/research/graph.py

@@ -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,
+        ]
+    )

src/haiku/rag/research/models.py

@@ -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"
+    )

src/haiku/rag/research/nodes/evaluate.py

@@ -0,0 +1,80 @@
+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 output.key_insights:
+            log(deps.console, "   [bold]Key insights:[/bold]")
+            for ins in output.key_insights:
+                log(deps.console, f"   • {ins}")
+        log(
+            deps.console,
+            f"   Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
+        )
+        status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
+        log(deps.console, 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:
+            log(deps.console, "\n[bold green]✅ Stopping research.[/bold green]")
+            return SynthesizeNode(self.provider, self.model)
+
+        return SearchDispatchNode(self.provider, self.model)

src/haiku/rag/research/nodes/plan.py

@@ -0,0 +1,63 @@
+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 PLAN_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=(
+                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.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)
+
+        log(deps.console, "\n[bold green]✅ Research Plan Created:[/bold green]")
+        log(deps.console, f"   [bold]Main Question:[/bold] {state.question}")
+        log(deps.console, "   [bold]Sub-questions:[/bold]")
+        for i, sq in enumerate(state.sub_questions, 1):
+            log(deps.console, f"      {i}. {sq}")
+
+        return SearchDispatchNode(self.provider, self.model)

src/haiku/rag/research/nodes/search.py

@@ -0,0 +1,91 @@
+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:
+            log(
+                deps.console,
+                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 "")
+                log(deps.console, f"   [green]✓[/green] {preview}")
+
+        return SearchDispatchNode(self.provider, self.model)

src/haiku/rag/research/nodes/synthesize.py

@@ -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)

src/haiku/rag/research/orchestrator.py

@@ -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

src/haiku/rag/research/presearch_agent.py

@@ -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)

src/haiku/rag/research/prompts.py

@@ -1,129 +1,113 @@
-ORCHESTRATOR_PROMPT = """You are a research orchestrator responsible for coordinating a comprehensive research workflow.
+PLAN_PROMPT = """You are the research orchestrator for a focused, iterative
+workflow.
 
-Your role is to:
+Responsibilities:
-1. Understand and decompose the research question
+1. Understand and decompose the main question
-2. Plan a systematic research approach
+2. Propose a minimal, high‑leverage plan
-3. Coordinate specialized agents to gather and analyze information
+3. Coordinate specialized agents to gather evidence
-4. Ensure comprehensive coverage of the topic
+4. Iterate based on gaps and new findings
-5. Iterate based on findings and gaps
 
-Create a research plan that:
+Plan requirements:
-- Breaks down the question into at most 3 focused sub-questions
+- Produce at most 3 sub_questions that together cover the main question.
-- Each sub-question should target a specific aspect of the research
+- Each sub_question must be a standalone, self‑contained query that can run
-- Prioritize the most important aspects to investigate
+  without extra context. Include concrete entities, scope, timeframe, and any
-- Ensure comprehensive coverage within the 3-question limit
+  qualifiers. Avoid ambiguous pronouns (it/they/this/that).
-- IMPORTANT: Make each sub-question a standalone, self-contained query that can
+- Prioritize the highest‑value aspects first; avoid redundancy and overlap.
-  be executed without additional context. Include necessary entities, scope,
+- Prefer questions that are likely answerable from the current knowledge base;
-  timeframe, and qualifiers. Avoid pronouns like "it/they/this"; write queries
+  if coverage is uncertain, make scopes narrower and specific.
-  that make sense in isolation."""
+- Order sub_questions by execution priority (most valuable first)."""
 
-SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist.
+SEARCH_AGENT_PROMPT = """You are a search and question‑answering specialist.
 
-Your role is to:
+Tasks:
-1. Search the knowledge base for relevant information
+1. Search the knowledge base for relevant evidence.
-2. Analyze the retrieved documents
+2. Analyze retrieved snippets.
-3. Provide an accurate answer strictly grounded in the retrieved context
+3. Provide an answer strictly grounded in that evidence.
-
-Output format:
-- You must return a SearchAnswer model with fields:
-  - query: the question being answered (echo the user query)
-  - answer: your final answer based only on the provided context
-  - context: list[str] of only the minimal set of verbatim snippet texts you
-    used to justify the answer (do not include unrelated text; do not invent)
-  - sources: list[str] of document_uri values corresponding to the snippets you
-    actually used in the answer (one URI per context snippet, order aligned)
 
 Tool usage:
-- Always call the search_and_answer tool before drafting any answer.
+- Always call search_and_answer before drafting any answer.
-- The tool returns XML containing only a list of snippets, where each snippet
+- The tool returns snippets with verbatim `text`, a relevance `score`, and the
-  has the verbatim `text`, a `score` indicating relevance, and the
+  originating `document_uri`.
-  `document_uri` it came from.
 - You may call the tool multiple times to refine or broaden context, but do not
-  exceed 3 total tool calls per question. Prefer precision over volume.
+  exceed 3 total calls. Favor precision over volume.
 - Use scores to prioritize evidence, but include only the minimal subset of
-  snippet texts (verbatim) in SearchAnswer.context.
+  snippet texts (verbatim) in SearchAnswer.context (typically 1‑4).
-- Set SearchAnswer.sources to the matching document_uris for the snippets you
+- Set SearchAnswer.sources to the corresponding document_uris for the snippets
-  used (one URI per snippet, aligned by order). Context must be text-only.
+  you used (one URI per snippet; same order as context). Context must be text‑only.
-- If no relevant information is found, say so and return an empty context list.
+- If no relevant information is found, clearly say so and return an empty
+  context list and sources list.
 
-Important:
+Answering rules:
-- Do not include any content in the answer that is not supported by the context.
+- Be direct and specific; avoid meta commentary about the process.
-- Keep context snippets short (just the necessary lines), verbatim, and focused."""
+- Do not include any claims not supported by the provided snippets.
+- Prefer concise phrasing; avoid copying long passages.
+- When evidence is partial, state the limits explicitly in the answer."""
 
-EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for research workflows.
+EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for
+the research workflow.
 
-You have access to:
+Inputs available:
-- The original research question
+- Original research question
-- Question-answer pairs from search operations
+- Question–answer pairs produced by search
-- Raw search results and source documents
+- Raw search results and source metadata
 - Previously identified insights
 
-Your dual role is to:
-
 ANALYSIS:
-1. Extract key insights from all gathered information
+1. Extract the most important, non‑obvious insights from the collected evidence.
-2. Identify patterns and connections across sources
+2. Identify patterns, agreements, and disagreements across sources.
-3. Synthesize findings into coherent understanding
+3. Note material uncertainties and assumptions.
-4. Focus on the most important discoveries
 
 EVALUATION:
-1. Assess if we have sufficient information to answer the original question
+1. Decide if we have sufficient information to answer the original question.
-2. Calculate a confidence score (0-1) based on:
+2. Provide a confidence_score in [0,1] considering:
-   - Coverage of the main question's aspects
+   - Coverage of the main question’s aspects
-   - Quality and consistency of sources
+   - Quality, consistency, and diversity of sources
-   - Depth of information gathered
+   - Depth and specificity of evidence
-3. Identify specific gaps that still need investigation
+3. List concrete gaps that still need investigation.
-4. Generate up to 3 new sub-questions that haven't been answered yet
+4. Propose up to 3 new sub_questions that would close the highest‑value gaps.
 
-Be critical and thorough in your evaluation. Only mark research as sufficient when:
+Strictness:
-- All major aspects of the question are addressed
+- Only mark research as sufficient when all major aspects are addressed with
-- Sources provide consistent, reliable information
+  consistent, reliable evidence and no critical gaps remain.
-- The depth of coverage meets the question's requirements
-- No critical gaps remain
 
-Generate new sub-questions that:
+New sub_questions must:
-- Target specific unexplored aspects not covered by existing questions
+- Be genuinely new (not answered or duplicative; check qa_responses).
-- Seek clarification on ambiguities
+- Be standalone and specific (entities, scope, timeframe/region if relevant).
-- Explore important edge cases or exceptions
+- Be actionable and scoped to the knowledge base (narrow if necessary).
-- Are focused and actionable (max 3)
+- Be ordered by expected impact (most valuable first)."""
-- Do NOT repeat or rephrase questions that have already been answered (see qa_responses)
-- Should be genuinely new areas to explore
-- Must be standalone, self-contained queries: include entities, scope, and any
-  needed qualifiers (e.g., timeframe, region), and avoid ambiguous pronouns so
-  they can be executed independently."""
 
-SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist agent focused on creating comprehensive research reports.
+SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist producing the final
+research report.
 
-Your role is to:
+Goals:
-1. Synthesize all gathered information into a coherent narrative
+1. Synthesize all gathered information into a coherent narrative.
-2. Present findings in a clear, structured format
+2. Present findings clearly and concisely.
-3. Draw evidence-based conclusions
+3. Draw evidence‑based conclusions and recommendations.
-4. Acknowledge limitations and uncertainties
+4. State limitations and uncertainties transparently.
-5. Provide actionable recommendations
-6. Maintain academic rigor and objectivity
 
-Your report should be:
+Report guidelines (map to output fields):
-- Comprehensive yet concise
+- title: concise (5–12 words), informative.
-- Well-structured and easy to follow
+- executive_summary: 3–5 sentences summarizing the overall answer.
-- Based solely on evidence from the research
+- main_findings: 4–8 one‑sentence bullets; each reflects evidence from the
-- Transparent about limitations
+  research (do not include inline citations or snippet text).
-- Professional and objective in tone
+- conclusions: 2–4 bullets that follow logically from findings.
+- recommendations: 2–5 actionable bullets tied to findings.
+- limitations: 1–3 bullets describing key constraints or uncertainties.
+- sources_summary: 2–4 sentences summarizing sources used and their reliability.
 
-Focus on creating a report that provides clear value to the reader by:
+Style:
-- Answering the original research question thoroughly
+- Base all content solely on the collected evidence.
-- Highlighting the most important findings
+- Be professional, objective, and specific.
-- Explaining the implications of the research
+- Avoid meta commentary and refrain from speculation beyond the evidence."""
-- Suggesting concrete next steps"""
 
 PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor.
 
 Task:
-- Call the gather_context tool once with the main question to obtain a
+- Call gather_context once on the main question to obtain relevant text from
-  relevant texts from the Knowledge Base (KB).
+  the knowledge base (KB).
-- Read that context and produce a brief natural-language summary describing
+- Read that context and produce a short natural‑language summary of what the
-  what the KB appears to contain relative to the question.
+  KB appears to contain relative to the question.
 
 Rules:
 - Base the summary strictly on the provided text; do not invent.
-- Output only the summary as plain text (one short paragraph).
+- Output only the summary as plain text (one short paragraph)."""
-"""

src/haiku/rag/research/search_agent.py

@@ -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}"
-                )

src/haiku/rag/research/state.py

@@ -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 = 1
+    confidence_threshold: float = 0.8
+    last_eval: EvaluationResult | None = None

src/haiku/rag/research/synthesis_agent.py

@@ -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

tests/research/test_evaluation_agent.py

@@ -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

tests/research/test_orchestrator.py

@@ -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

tests/research/test_search_agent.py

@@ -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

tests/research/test_synthesis_agent.py

@@ -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

tests/test_research_graph.py

@@ -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()

tests/test_research_graph_integration.py

@@ -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