Merge pull request #267 from ggozad/feat/research-graph-improvements

Performance improvements for the research graph & chat agents

CHANGELOG.md

@@ -1,6 +1,21 @@
 # Changelog
 ## [Unreleased]
 
+### Changed
+
+- **Iterative Research Planning**: Research graph now uses an iterative feedback loop instead of batch question processing
+  - Planner proposes ONE question at a time, sees the answer, then decides whether to continue
+  - Removes `gather_context` tool — planner proposes questions directly
+  - Simpler flow: `plan_next` → `search_one` → loop back until complete → `synthesize`
+  - Consolidated `build_conversational_graph()` into `build_research_graph(output_mode="conversational")`
+
+### Removed
+
+- **Dead config options**: Removed vestigial fields from iterative planning refactor
+  - `confidence_threshold` from `ResearchConfig` and `ResearchState` (LLM decides completion via `is_complete`)
+  - `max_sub_questions` from `QAConfig` (iterative flow uses one question at a time)
+  - `sub_questions` field from `ResearchContext` (no longer populated)
+
 ## [0.27.2] - 2026-01-29
 
 ### Added

docs/agents.md

@@ -163,44 +163,39 @@ Frontend clients should extract state from under this key. See the [Web Applicat
 
 ## Research Graph
 
-The research workflow is implemented as a typed pydantic-graph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report.
+The research workflow is implemented as a typed pydantic-graph. It uses an iterative feedback loop where the planner proposes one question at a time, sees the answer, then decides whether to continue or synthesize.
 
 ```mermaid
 ---
 title: Research graph
 ---
 stateDiagram-v2
-  [*] --> plan
+  [*] --> plan_next
-  plan --> get_batch
+  plan_next --> search_one: Has next question
-  get_batch --> search_one: Has questions (map)
+  plan_next --> synthesize: Complete or max iterations
-  get_batch --> synthesize: No questions
+  search_one --> plan_next
-  search_one --> collect_answers
-  collect_answers --> decide
-  decide --> get_batch: Continue research
-  decide --> synthesize: Done researching
   synthesize --> [*]
 ```
 
 **Key nodes:**
 
-- **plan**: Builds up to 3 standalone sub-questions (uses an internal presearch tool)
+- **plan_next**: Evaluates gathered evidence and either proposes the next question to investigate or marks research as complete
-- **get_batch**: Retrieves remaining sub-questions for the current iteration
+- **search_one**: Answers a single question using the knowledge base
-- **search_one**: Answers a single sub-question using the KB (mapped in parallel)
-- **collect_answers**: Aggregates search results from parallel executions
-- **decide**: Evaluates confidence and determines whether to continue or synthesize
 - **synthesize**: Generates a final structured research report
 
 **Primary models:**
 
-- `SearchAnswer` — one per sub-question (query, answer, confidence, citations)
+- `IterativePlanResult` — planning decision (is_complete, next_question, reasoning)
-- `EvaluationResult` — confidence score, new questions, sufficiency assessment
+- `SearchAnswer` — answer to a single question (query, answer, confidence, citations)
 - `ResearchReport` — final report (title, executive summary, findings, conclusions, …)
+- `ConversationalAnswer` — alternative output for chat integration (answer, citations, confidence)
 
-**Parallel execution:**
+**Iterative flow:**
 
-- The `search_one` node is mapped over all questions in a batch
+- Each iteration: planner evaluates context → proposes one question → search answers it → loop back
-- Parallelism is controlled via `max_concurrency`
+- Planner can decompose complex questions (e.g., "benefits and drawbacks" → start with "benefits")
-- Decision nodes process results after each batch completes
+- Session context is used to resolve ambiguous references and inform planning
+- Loop terminates when planner marks `is_complete=True` or `max_iterations` is reached
 
 ### CLI Usage
 
@@ -249,7 +244,6 @@ custom_config = AppConfig(
         provider="openai",
         model="gpt-4o-mini",
         max_iterations=5,
-        confidence_threshold=0.85,
         max_concurrency=3,
     )
 )

docs/cli.md

@@ -255,7 +255,7 @@ Flags:
 - `--context`: Background context for the research
 - `--context-file`: Path to a file containing background context
 
-Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.
+Research parameters like `max_iterations` and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.
 
 ## Server
 

docs/configuration/index.md

@@ -85,7 +85,6 @@ qa:
     provider: ollama
     name: gpt-oss
     enable_thinking: false
-  max_sub_questions: 3
   max_iterations: 2
   max_concurrency: 1
 
@@ -95,7 +94,6 @@ research:
     name: ""
     enable_thinking: false
   max_iterations: 3
-  confidence_threshold: 0.8
   max_concurrency: 1
 
 search:

docs/configuration/qa-research.md

@@ -32,17 +32,15 @@ qa:
     provider: ollama
     name: gpt-oss
     enable_thinking: false
-  max_sub_questions: 3  # Maximum sub-questions for deep QA
+  max_iterations: 2     # Maximum search iterations
-  max_iterations: 2     # Maximum search iterations per sub-question
+  max_concurrency: 1    # Concurrent search operations
-  max_concurrency: 1    # Sub-questions processed in parallel
 ```
 
 - **model**: LLM configuration (see [Providers](providers.md#model-settings))
-- **max_sub_questions**: For deep QA mode, maximum number of sub-questions to generate (default: 3)
+- **max_iterations**: Maximum search iterations (default: 2)
-- **max_iterations**: Maximum search/evaluate cycles per sub-question (default: 2)
+- **max_concurrency**: Number of concurrent search operations (default: 1)
-- **max_concurrency**: Number of sub-questions to process in parallel (default: 1)
 
-Deep QA mode (`haiku-rag ask --deep`) decomposes complex questions into sub-questions, processes them in parallel batches, and synthesizes the results.
+Deep QA mode (`haiku-rag ask --deep`) uses the research graph with a single iteration for quick, focused answers.
 
 ## Research Configuration
 
@@ -55,13 +53,11 @@ research:
     name: ""               # Empty to use qa model
     enable_thinking: false
   max_iterations: 3
-  confidence_threshold: 0.8
   max_concurrency: 1
 ```
 
 - **model**: LLM configuration. Leave provider/model empty to inherit from `qa` (see [Providers](providers.md#model-settings))
-- **max_iterations**: Maximum search/evaluate cycles (default: 3)
+- **max_iterations**: Maximum planning/search iterations (default: 3)
-- **confidence_threshold**: Stop when confidence score meets/exceeds this (default: 0.8)
+- **max_concurrency**: Concurrent search operations (default: 1)
-- **max_concurrency**: Sub-questions searched in parallel per iteration (default: 1)
 
-The research workflow plans sub-questions, searches in parallel batches, evaluates findings, and iterates until reaching the confidence threshold or max iterations.
+The research workflow uses an iterative feedback loop: the planner proposes one question at a time, sees the answer, then decides whether to continue or synthesize. This continues until the planner marks research as complete or `max_iterations` is reached.

evaluations/evaluations/benchmark.py

@@ -308,12 +308,7 @@ async def run_qa_benchmark(
 
             async def answer_question(question: str) -> str:
                 context = ResearchContext(original_question=question)
-                state = ResearchState.from_config(
+                state = ResearchState.from_config(context=context, config=config)
-                    context=context,
-                    config=config,
-                    max_iterations=2,
-                    confidence_threshold=0.0,
-                )
                 deps = ResearchDeps(client=rag)
                 report = await graph.run(state=state, deps=deps)
                 return report.executive_summary if report else ""

haiku_rag_slim/haiku/rag/agents/__init__.py

@@ -9,16 +9,13 @@ from haiku.rag.agents.chat import (
 from haiku.rag.agents.qa import QuestionAnswerAgent, get_qa_agent
 from haiku.rag.agents.research import (
     Citation,
-    EvaluationResult,
+    IterativePlanResult,
     ResearchContext,
     ResearchDependencies,
     ResearchReport,
     SearchAnswer,
 )
-from haiku.rag.agents.research.graph import (
+from haiku.rag.agents.research.graph import build_research_graph
-    build_conversational_graph,
-    build_research_graph,
-)
 from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 
 __all__ = [
@@ -27,7 +24,6 @@ __all__ = [
     "QuestionAnswerAgent",
     # Research
     "build_research_graph",
-    "build_conversational_graph",
     "ResearchContext",
     "ResearchDependencies",
     "ResearchDeps",
@@ -35,7 +31,7 @@ __all__ = [
     "ResearchReport",
     "Citation",
     "SearchAnswer",
-    "EvaluationResult",
+    "IterativePlanResult",
     # Chat
     "create_chat_agent",
     "SearchAgent",

haiku_rag_slim/haiku/rag/agents/chat/agent.py

@@ -23,7 +23,7 @@ from haiku.rag.agents.chat.state import (
     emit_state_event,
 )
 from haiku.rag.agents.research.dependencies import ResearchContext
-from haiku.rag.agents.research.graph import build_conversational_graph
+from haiku.rag.agents.research.graph import build_research_graph
 from haiku.rag.agents.research.models import Citation
 from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 from haiku.rag.client import HaikuRAG
@@ -197,7 +197,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         doc_filter = combine_filters(session_filter, tool_filter)
 
         # Build and run the conversational research graph
-        graph = build_conversational_graph(config=ctx.deps.config)
+        graph = build_research_graph(
+            config=ctx.deps.config, output_mode="conversational"
+        )
         session_id = ctx.deps.session_state.session_id
 
         # Get session context from server cache for planning, fallback to initial_context
@@ -247,7 +249,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         state = ResearchState(
             context=context,
             max_iterations=1,
-            confidence_threshold=0.0,
             search_filter=doc_filter,
             max_concurrency=ctx.deps.config.research.max_concurrency,
         )

haiku_rag_slim/haiku/rag/agents/research/__init__.py

@@ -1,7 +1,7 @@
 from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
 from haiku.rag.agents.research.models import (
     Citation,
-    EvaluationResult,
+    IterativePlanResult,
     ResearchReport,
     SearchAnswer,
 )

haiku_rag_slim/haiku/rag/agents/research/dependencies.py

@@ -13,9 +13,6 @@ class ResearchContext(BaseModel):
     """Context shared across research agents."""
 
     original_question: str = Field(description="The original research question")
-    sub_questions: list[str] = Field(
-        default_factory=list, description="Decomposed sub-questions"
-    )
     qa_responses: list[Any] = Field(
         default_factory=list, description="Structured QA pairs used during research"
     )

haiku_rag_slim/haiku/rag/agents/research/graph.py

@@ -1,25 +1,23 @@
 import asyncio
+from typing import Literal, overload
 
 from pydantic_ai import Agent, RunContext, format_as_xml
 from pydantic_ai.output import ToolOutput
 from pydantic_graph.beta import Graph, GraphBuilder, StepContext
-from pydantic_graph.beta.join import reduce_list_append
 
 from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
 from haiku.rag.agents.research.models import (
     Citation,
     ConversationalAnswer,
-    EvaluationResult,
+    IterativePlanResult,
     RawSearchAnswer,
-    ResearchPlan,
     ResearchReport,
     SearchAnswer,
 )
 from haiku.rag.agents.research.prompts import (
     CONVERSATIONAL_SYNTHESIS_PROMPT,
-    DECISION_PROMPT,
+    ITERATIVE_PLAN_PROMPT,
-    PLAN_PROMPT,
+    ITERATIVE_PLAN_PROMPT_WITH_CONTEXT,
-    PLAN_PROMPT_WITH_CONTEXT,
     SEARCH_PROMPT,
     SYNTHESIS_PROMPT,
 )
@@ -29,17 +27,8 @@ from haiku.rag.config.models import AppConfig
 from haiku.rag.utils import build_prompt, get_model
 
 
-def format_context_for_prompt(
+def format_context_for_prompt(context: ResearchContext) -> str:
-    context: ResearchContext,
+    """Format the research context as XML for prompts."""
-    include_pending_questions: bool = True,
-) -> str:
-    """Format the research context as XML for prompts.
-
-    Args:
-        context: The research context to format.
-        include_pending_questions: Whether to include pending sub-questions.
-            Set to False for synthesis prompts where pending questions aren't relevant.
-    """
     context_data: dict[str, object] = {}
 
     if context.session_context:
@@ -47,9 +36,6 @@ def format_context_for_prompt(
 
     context_data["question"] = context.original_question
 
-    if include_pending_questions and context.sub_questions:
-        context_data["pending_questions"] = context.sub_questions
-
     if context.qa_responses:
         context_data["prior_answers"] = [
             {
@@ -64,80 +50,63 @@ def format_context_for_prompt(
     return format_as_xml(context_data, root_tag="context")
 
 
-# =============================================================================
+async def _iterative_plan_logic(
-# Shared step logic helpers
-# =============================================================================
-
-
-async def _plan_step_logic(
     state: ResearchState,
     deps: ResearchDeps,
     config: AppConfig,
-    plan_prompt: str,
+) -> IterativePlanResult:
-) -> None:
+    """Evaluate context and decide next question or mark complete."""
-    """Shared logic for the plan step."""
+    has_prior_answers = bool(state.context.qa_responses)
+
+    # If max iterations reached, skip LLM and mark complete
+    if state.iterations >= state.max_iterations:
+        return IterativePlanResult(
+            is_complete=True,
+            next_question=None,
+            reasoning=f"Max iterations ({state.max_iterations}) reached.",
+        )
+
     model_config = config.research.model
 
-    # Use context-aware prompt if we have existing qa_responses or session_context
+    if has_prior_answers:
-    has_prior_answers = bool(state.context.qa_responses)
+        effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT_WITH_CONTEXT, config)
-    has_session_context = bool(state.context.session_context)
+    else:
-    effective_plan_prompt = (
+        effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT, config)
-        build_prompt(PLAN_PROMPT_WITH_CONTEXT, config)
-        if has_prior_answers or has_session_context
-        else plan_prompt
-    )
 
-    plan_agent: Agent[ResearchDependencies, ResearchPlan] = Agent(  # type: ignore[invalid-assignment]
+    plan_agent: Agent[ResearchDependencies, IterativePlanResult] = Agent(  # type: ignore[assignment]
         model=get_model(model_config, config),
-        output_type=ResearchPlan,
+        output_type=IterativePlanResult,
-        instructions=effective_plan_prompt,
+        instructions=effective_prompt,
         retries=3,
         output_retries=3,
         deps_type=ResearchDependencies,
     )
 
-    search_filter = state.search_filter
+    # Build prompt based on current state
-
-    # Only register gather_context tool when we don't have existing context
-    if not has_prior_answers and not has_session_context:
-
-        @plan_agent.tool
-        async def gather_context(
-            ctx2: RunContext[ResearchDependencies],
-            query: str,
-            limit: int | None = None,
-        ) -> str:
-            results = await ctx2.deps.client.search(
-                query, limit=limit, filter=search_filter
-            )
-            results = await ctx2.deps.client.expand_context(results)
-            return "\n\n".join(r.content for r in results)
-
-    # Build prompt with existing context if available
     if has_prior_answers:
         context_xml = format_context_for_prompt(state.context)
         prompt = (
-            f"Review existing context and plan additional research if needed.\n\n"
+            f"Review the gathered evidence and decide whether to continue or synthesize.\n\n"
-            f"{context_xml}\n\n"
+            f"{context_xml}"
-            f"Main question: {state.context.original_question}"
-        )
-    elif has_session_context:
-        context_xml = format_context_for_prompt(state.context)
-        prompt = (
-            f"Plan a focused approach for the main question.\n\n"
-            f"{context_xml}\n\n"
-            f"Main question: {state.context.original_question}"
         )
     else:
-        prompt = (
+        context_xml = format_context_for_prompt(state.context)
-            "Plan a focused approach for the main question.\n\n"
+        prompt = f"Plan the research investigation.\n\n{context_xml}"
-            f"Main question: {state.context.original_question}"
-        )
 
     agent_deps = ResearchDependencies(client=deps.client, context=state.context)
-    plan_result = await plan_agent.run(prompt, deps=agent_deps)
+    result = await plan_agent.run(prompt, deps=agent_deps)
-    output = plan_result.output
+
-    state.context.sub_questions = list(output.sub_questions)
+    # Enforce: if no prior answers, must have a next_question to investigate
+    if not has_prior_answers:
+        if result.output.is_complete or not result.output.next_question:
+            return IterativePlanResult(
+                is_complete=False,
+                next_question=result.output.next_question
+                or state.context.original_question,
+                reasoning=result.output.reasoning,
+            )
+
+    return result.output
 
 
 async def _search_one_step_logic(
@@ -147,14 +116,14 @@ async def _search_one_step_logic(
     search_prompt: str,
     sub_q: str,
 ) -> SearchAnswer:
-    """Shared logic for the search_one step."""
+    """Answer a single question using the knowledge base."""
     model_config = config.research.model
 
     if deps.semaphore is None:
         deps.semaphore = asyncio.Semaphore(state.max_concurrency)
 
     async with deps.semaphore:
-        agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent(  # type: ignore[invalid-assignment]
+        agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent(  # type: ignore[assignment]
             model=get_model(model_config, config),
             output_type=ToolOutput(RawSearchAnswer, max_retries=3),
             instructions=search_prompt,
@@ -176,7 +145,6 @@ async def _search_one_step_logic(
             )
             results = await ctx2.deps.client.expand_context(results)
             ctx2.deps.search_results = results
-            # Format with rank instead of raw score to avoid confusing LLMs
             total = len(results)
             parts = [
                 r.format_for_agent(rank=i + 1, total=total)
@@ -190,6 +158,10 @@ async def _search_one_step_logic(
 
         result = await agent.run(sub_q, deps=agent_deps)
         raw_answer = result.output
+
+        # Increment iterations after each search completes
+        state.iterations += 1
+
         if raw_answer:
             answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
             state.context.add_qa_response(answer)
@@ -197,262 +169,62 @@ async def _search_one_step_logic(
         return SearchAnswer(query=sub_q, answer="", confidence=0.0)
 
 
-def _get_batch_logic(state: ResearchState) -> list[str] | None:
+@overload
-    """Shared logic for the get_batch step."""
+def build_research_graph(
-    if not state.context.sub_questions:
+    config: AppConfig = ...,
-        return None
+    output_mode: Literal["report"] = ...,
-
+) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]: ...
-    batch = list(state.context.sub_questions)
-    state.context.sub_questions.clear()
-    return batch
 
 
-# =============================================================================
+@overload
-# Research graph (full version with decide loop)
+def build_research_graph(
-# =============================================================================
+    config: AppConfig = ...,
+    output_mode: Literal["conversational"] = ...,
+) -> Graph[ResearchState, ResearchDeps, None, ConversationalAnswer]: ...
 
 
 def build_research_graph(
     config: AppConfig = Config,
-    include_plan: bool = True,
+    output_mode: Literal["report", "conversational"] = "report",
-) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
+) -> Graph[ResearchState, ResearchDeps, None, ResearchReport | ConversationalAnswer]:
-    """Build the Research graph.
+    """Build the iterative research graph.
 
     Args:
         config: AppConfig object (uses config.research for provider, model, and graph parameters)
-        include_plan: Whether to include the planning step (False for execute-only mode)
+        output_mode: Output format - "report" for ResearchReport, "conversational" for ConversationalAnswer
 
     Returns:
-        Configured Research graph
+        Configured research graph with iterative planning
     """
     model_config = config.research.model
 
-    # Build prompts with system_context if configured
-    plan_prompt = build_prompt(
-        PLAN_PROMPT
-        + "\n\nUse the gather_context tool once on the main question before planning.",
-        config,
-    )
     search_prompt = build_prompt(SEARCH_PROMPT, config)
-    decision_prompt = build_prompt(DECISION_PROMPT, config)
-    synthesis_prompt = build_prompt(
-        config.prompts.synthesis or SYNTHESIS_PROMPT, config
-    )
-    g = GraphBuilder(
-        state_type=ResearchState,
-        deps_type=ResearchDeps,
-        output_type=ResearchReport,
-    )
 
-    @g.step
+    if output_mode == "report":
-    async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
+        synthesis_prompt = build_prompt(
-        """Create research plan with sub-questions."""
+            config.prompts.synthesis or SYNTHESIS_PROMPT, config
-        await _plan_step_logic(ctx.state, ctx.deps, config, plan_prompt)
-
-    @g.step
-    async def search_one(
-        ctx: StepContext[ResearchState, ResearchDeps, str],
-    ) -> SearchAnswer:
-        """Answer a single sub-question using the knowledge base."""
-        try:
-            return await _search_one_step_logic(
-                ctx.state, ctx.deps, config, search_prompt, ctx.inputs
-            )
-        except Exception as e:
-            return SearchAnswer(
-                query=ctx.inputs,
-                answer=f"Search failed: {str(e)}",
-                confidence=0.0,
-            )
-
-    @g.step
-    async def get_batch(
-        ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
-    ) -> list[str] | None:
-        """Get all remaining questions for this iteration."""
-        return _get_batch_logic(ctx.state)
-
-    @g.step
-    async def decide(
-        ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
-    ) -> bool:
-        """Evaluate research sufficiency and decide whether to continue."""
-        state = ctx.state
-        deps = ctx.deps
-
-        agent: Agent[ResearchDependencies, EvaluationResult] = Agent(  # type: ignore[invalid-assignment]
-            model=get_model(model_config, config),
-            output_type=EvaluationResult,
-            instructions=decision_prompt,
-            retries=3,
-            output_retries=3,
-            deps_type=ResearchDependencies,
-        )
-
-        context_xml = format_context_for_prompt(state.context)
-        prompt_parts = [
-            "Assess whether the research now answers the original question with adequate confidence.",
-            context_xml,
-        ]
-        if state.last_eval is not None:
-            prev = state.last_eval
-            prompt_parts.append(
-                "<previous_evaluation>"
-                f"<confidence>{prev.confidence_score:.2f}</confidence>"
-                f"<is_sufficient>{str(prev.is_sufficient).lower()}</is_sufficient>"
-                f"<reasoning>{prev.reasoning}</reasoning>"
-                "</previous_evaluation>"
-            )
-        prompt = "\n\n".join(part for part in prompt_parts if part)
-
-        agent_deps = ResearchDependencies(
-            client=deps.client,
-            context=state.context,
-        )
-        decision_result = await agent.run(prompt, deps=agent_deps)
-        output = decision_result.output
-
-        state.last_eval = output
-        state.iterations += 1
-
-        # Get already-answered questions to avoid duplicates
-        answered_queries = {qa.query.lower() for qa in state.context.qa_responses}
-
-        for new_q in output.new_questions:
-            # Skip if already in pending or already answered
-            if new_q in state.context.sub_questions:
-                continue
-            if new_q.lower() in answered_queries:
-                continue
-            state.context.sub_questions.append(new_q)
-
-        should_continue = (
-            not output.is_sufficient
-            or output.confidence_score < state.confidence_threshold
-        ) and state.iterations < state.max_iterations
-
-        return should_continue
-
-    @g.step
-    async def synthesize(
-        ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
-    ) -> ResearchReport:
-        """Generate final research report."""
-        state = ctx.state
-        deps = ctx.deps
-
-        agent: Agent[ResearchDependencies, ResearchReport] = Agent(  # type: ignore[invalid-assignment]
-            model=get_model(model_config, config),
-            output_type=ResearchReport,
-            instructions=synthesis_prompt,
-            retries=3,
-            output_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,
-        )
-        result = await agent.run(prompt, deps=agent_deps)
-        return result.output
-
-    # Build the graph structure
-    collect_answers = g.join(
-        reduce_list_append,
-        initial_factory=list[SearchAnswer],
-    )
-
-    if include_plan:
-        g.add(
-            g.edge_from(g.start_node).to(plan),
-            g.edge_from(plan).to(get_batch),
         )
     else:
-        g.add(g.edge_from(g.start_node).to(get_batch))
+        synthesis_prompt = build_prompt(CONVERSATIONAL_SYNTHESIS_PROMPT, config)
-
-    g.add(
-        g.edge_from(get_batch).to(
-            g.decision()
-            .branch(g.match(list).label("Has questions").map().to(search_one))
-            .branch(g.match(type(None)).label("No questions").to(synthesize))
-        ),
-        g.edge_from(search_one).to(collect_answers),
-        g.edge_from(collect_answers).to(decide),
-    )
-
-    g.add(
-        g.edge_from(decide).to(
-            g.decision()
-            .branch(
-                g.match(bool, matches=lambda x: x)
-                .label("Continue research")
-                .to(get_batch)
-            )
-            .branch(
-                g.match(bool, matches=lambda x: not x)
-                .label("Done researching")
-                .to(synthesize)
-            )
-        ),
-        g.edge_from(synthesize).to(g.end_node),
-    )
-
-    return g.build()
-
-
-# =============================================================================
-# Conversational graph (simplified, single iteration)
-# =============================================================================
-
-
-def build_conversational_graph(
-    config: AppConfig = Config,
-) -> Graph[ResearchState, ResearchDeps, None, ConversationalAnswer]:
-    """Build a simplified research graph for conversational chat.
-
-    This graph is optimized for single-iteration Q&A:
-    - Context-aware planning (generates fewer sub-questions when context exists)
-    - Single search iteration (no decide loop)
-    - Conversational output (direct answer, not formal report)
-
-    Args:
-        config: AppConfig object
-
-    Returns:
-        Graph that outputs ConversationalAnswer
-    """
-    # Build prompts
-    plan_prompt = build_prompt(
-        PLAN_PROMPT
-        + "\n\nUse the gather_context tool once on the main question before planning.",
-        config,
-    )
-    search_prompt = build_prompt(SEARCH_PROMPT, config)
-    conversational_prompt = build_prompt(CONVERSATIONAL_SYNTHESIS_PROMPT, config)
 
     g = GraphBuilder(
         state_type=ResearchState,
         deps_type=ResearchDeps,
-        output_type=ConversationalAnswer,
+        output_type=ResearchReport if output_mode == "report" else ConversationalAnswer,
     )
 
     @g.step
-    async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
+    async def plan_next(
-        """Create research plan with sub-questions."""
+        ctx: StepContext[ResearchState, ResearchDeps, None | SearchAnswer],
-        await _plan_step_logic(ctx.state, ctx.deps, config, plan_prompt)
+    ) -> IterativePlanResult:
+        """Evaluate context and decide next question or complete."""
+        return await _iterative_plan_logic(ctx.state, ctx.deps, config)
 
     @g.step
     async def search_one(
         ctx: StepContext[ResearchState, ResearchDeps, str],
     ) -> SearchAnswer:
-        """Answer a single sub-question using the knowledge base."""
+        """Answer a single question using the knowledge base."""
         try:
             return await _search_one_step_logic(
                 ctx.state, ctx.deps, config, search_prompt, ctx.inputs
@@ -464,71 +236,114 @@ def build_conversational_graph(
                 confidence=0.0,
             )
 
-    @g.step
+    if output_mode == "report":
-    async def get_batch(
-        ctx: StepContext[ResearchState, ResearchDeps, None],
-    ) -> list[str] | None:
-        """Get all remaining questions for this iteration."""
-        return _get_batch_logic(ctx.state)
 
-    @g.step
+        @g.step
-    async def synthesize(
+        async def synthesize(
-        ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer] | None],
+            ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
-    ) -> ConversationalAnswer:
+        ) -> ResearchReport:
-        """Generate conversational answer from gathered evidence."""
+            """Generate final research report."""
-        state = ctx.state
+            state = ctx.state
-        deps = ctx.deps
+            deps = ctx.deps
 
-        agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent(  # type: ignore[invalid-assignment]
+            agent: Agent[ResearchDependencies, ResearchReport] = Agent(  # type: ignore[assignment]
-            model=get_model(config.research.model, config),
+                model=get_model(model_config, config),
-            output_type=ConversationalAnswer,
+                output_type=ResearchReport,
-            instructions=conversational_prompt,
+                instructions=synthesis_prompt,
-            retries=3,
+                retries=3,
-            output_retries=3,
+                output_retries=3,
-            deps_type=ResearchDependencies,
+                deps_type=ResearchDependencies,
-        )
+            )
 
-        context_xml = format_context_for_prompt(
+            context_xml = format_context_for_prompt(state.context)
-            state.context, include_pending_questions=False
+            prompt = (
-        )
+                "Generate a comprehensive research report based on all gathered information.\n\n"
-        prompt = f"Answer the question based on the gathered evidence.\n\n{context_xml}"
+                f"{context_xml}\n\n"
-        agent_deps = ResearchDependencies(
+                "Create a detailed report that synthesizes all findings into a coherent response."
-            client=deps.client,
+            )
-            context=state.context,
+            agent_deps = ResearchDependencies(
-        )
+                client=deps.client,
-        result = await agent.run(prompt, deps=agent_deps)
+                context=state.context,
+            )
+            result = await agent.run(prompt, deps=agent_deps)
+            return result.output
 
-        # Collect unique citations from qa_responses (dedupe by chunk_id)
+    else:
-        seen_chunks: set[str] = set()
-        unique_citations: list[Citation] = []
-        for qa in state.context.qa_responses:
-            for c in qa.citations:
-                if c.chunk_id not in seen_chunks:
-                    seen_chunks.add(c.chunk_id)
-                    unique_citations.append(c)
 
-        return ConversationalAnswer(
+        @g.step
-            answer=result.output.answer,
+        async def synthesize(
-            citations=unique_citations,
+            ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
-            confidence=result.output.confidence,
+        ) -> ConversationalAnswer:
-        )
+            """Generate conversational answer from gathered evidence."""
+            state = ctx.state
+            deps = ctx.deps
 
-    # Build the graph structure (simplified: plan → search → synthesize)
+            agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent(  # type: ignore[assignment]
-    collect_answers = g.join(
+                model=get_model(model_config, config),
-        reduce_list_append,
+                output_type=ConversationalAnswer,
-        initial_factory=list[SearchAnswer],
+                instructions=synthesis_prompt,
-    )
+                retries=3,
+                output_retries=3,
+                deps_type=ResearchDependencies,
+            )
+
+            context_xml = format_context_for_prompt(state.context)
+            prompt = (
+                f"Answer the question based on the gathered evidence.\n\n{context_xml}"
+            )
+            agent_deps = ResearchDependencies(
+                client=deps.client,
+                context=state.context,
+            )
+            result = await agent.run(prompt, deps=agent_deps)
+
+            # Collect unique citations from qa_responses (dedupe by chunk_id)
+            seen_chunks: set[str] = set()
+            unique_citations: list[Citation] = []
+            for qa in state.context.qa_responses:
+                for c in qa.citations:
+                    if c.chunk_id not in seen_chunks:
+                        seen_chunks.add(c.chunk_id)
+                        unique_citations.append(c)
+
+            return ConversationalAnswer(
+                answer=result.output.answer,
+                citations=unique_citations,
+                confidence=result.output.confidence,
+            )
+
+    # Build graph edges: iterative loop
+    #
+    # START -> plan_next -> [decision]
+    #                          |
+    #          [is_complete or max_iterations] -> synthesize -> END
+    #                          |
+    #          [has next_question] -> search_one -> plan_next (loop)
+
+    def extract_question(
+        ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
+    ) -> str:
+        """Extract next_question from IterativePlanResult."""
+        return ctx.inputs.next_question or ""
 
     g.add(
-        g.edge_from(g.start_node).to(plan),
+        g.edge_from(g.start_node).to(plan_next),
-        g.edge_from(plan).to(get_batch),
+        g.edge_from(plan_next).to(
-        g.edge_from(get_batch).to(
             g.decision()
-            .branch(g.match(list).label("Has questions").map().to(search_one))
+            .branch(
-            .branch(g.match(type(None)).label("No questions").to(synthesize))
+                g.match(
+                    IterativePlanResult,
+                    matches=lambda r: not r.is_complete and r.next_question is not None,
+                )
+                .label("Continue research")
+                .transform(extract_question)
+                .to(search_one)
+            )
+            .branch(
+                g.match(IterativePlanResult).label("Done researching").to(synthesize)
+            )
         ),
-        g.edge_from(search_one).to(collect_answers),
+        g.edge_from(search_one).to(plan_next),
-        g.edge_from(collect_answers).to(synthesize),
         g.edge_from(synthesize).to(g.end_node),
     )
 

haiku_rag_slim/haiku/rag/agents/research/models.py

@@ -1,25 +1,21 @@
 from typing import TYPE_CHECKING
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, Field
 
 if TYPE_CHECKING:
     from haiku.rag.store.models import SearchResult
 
 
-class ResearchPlan(BaseModel):
+class IterativePlanResult(BaseModel):
-    """A structured research plan with sub-questions to explore."""
+    """Output from iterative planning step."""
 
-    sub_questions: list[str] = Field(
+    is_complete: bool = Field(
-        ...,
+        description="Whether research is complete and can be synthesized"
-        description="Specific questions to research, phrased as complete questions",
     )
-
+    next_question: str | None = Field(
-    @field_validator("sub_questions")
+        default=None, description="Next question to investigate, if not complete"
-    @classmethod
+    )
-    def validate_sub_questions(cls, v: list[str]) -> list[str]:
+    reasoning: str = Field(description="Brief explanation of the decision")
-        if len(v) > 12:
-            raise ValueError("Cannot have more than 12 sub-questions")
-        return v
 
 
 class Citation(BaseModel):
@@ -115,27 +111,6 @@ def resolve_citations(
     return citations
 
 
-class EvaluationResult(BaseModel):
-    """Result of research sufficiency evaluation."""
-
-    is_sufficient: bool = Field(
-        description="Whether the research is sufficient to answer the original question"
-    )
-    confidence_score: float = Field(
-        ge=0.0,
-        le=1.0,
-        description="Confidence level in the completeness of research (0-1)",
-    )
-    reasoning: str = Field(
-        description="Explanation of why the research is or isn't complete"
-    )
-    new_questions: list[str] = Field(
-        default_factory=list,
-        max_length=3,
-        description="New sub-questions to add to the research (max 3)",
-    )
-
-
 class ConversationalAnswer(BaseModel):
     """Conversational answer for chat context."""
 

haiku_rag_slim/haiku/rag/agents/research/prompts.py

@@ -1,47 +1,47 @@
-PLAN_PROMPT = """You are the research orchestrator for a focused workflow.
+ITERATIVE_PLAN_PROMPT = """You are the research orchestrator planning the investigation.
 
-If a <background> section is provided, use it to understand the domain context.
+If a <background> section is provided, use it to understand the conversation context.
 
-Responsibilities:
+Your task:
-1. Understand and decompose the main question
+1. Analyze the original question
-2. Propose a minimal, high-leverage plan
+2. Propose the first question to investigate
-3. Coordinate specialized agents to gather evidence
 
-Plan requirements:
+For simple questions, investigate them directly. For composite or complex questions,
-- Produce at most 3 sub_questions that together cover the main question.
+you may decompose into a focused sub-question. For example:
-- sub_questions must be a list of plain strings, where each string is a complete
+- "What are the benefits and drawbacks of X?" → Start with "What are the benefits of X?"
-  question. Do NOT use objects with nested fields like {question, details}.
+- Ambiguous references should be resolved using background context if available
-- Each sub_question must be a standalone, self-contained query that can run
-  without extra context. Include concrete entities, scope, timeframe, and any
-  qualifiers. Avoid ambiguous pronouns (it/they/this/that).
-- Prioritize the highest-value aspects first; avoid redundancy and overlap.
-- Prefer questions that are likely answerable from the current knowledge base;
-  if coverage is uncertain, make scopes narrower and specific.
-- Order sub_questions by execution priority (most valuable first).
 
-Use the gather_context tool once on the main question before planning."""
+Output requirements:
+- Set is_complete=False (you are just starting the investigation)
+- Set next_question to the question to investigate
+- Provide brief reasoning explaining your choice
 
-PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator for a focused workflow.
+The question must be standalone and self-contained:
+- Include concrete entities, scope, and any qualifiers
+- Avoid ambiguous pronouns (it/they/this/that)"""
+
+ITERATIVE_PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator evaluating gathered evidence.
 
 You have access to context that may include:
 - <background>: Domain context for the conversation
 - <prior_answers>: Previous Q&A pairs with confidence scores
 
-Review the provided context first. Use <background> to understand the domain.
+Your task:
-If <prior_answers> exist and already answer the question completely,
+1. Review the provided evidence carefully
-return an empty sub_questions list. Only create sub-questions to fill gaps.
+2. Assess whether it sufficiently answers the original question
+3. Decide whether to continue research or synthesize
 
-Responsibilities:
+Decision criteria:
-1. Review provided context to understand what's already known
+- Set is_complete=True if the evidence adequately answers the question
-2. Identify gaps that need additional research
+- Set is_complete=False with a next_question if important gaps remain
-3. Propose minimal sub-questions only for missing information
 
-Plan requirements:
+If not complete, propose exactly ONE high-value follow-up question in next_question:
-- If existing context fully answers the question, return an empty sub_questions list.
+- Focus on the most critical gap not covered by prior_answers
-- Only create new sub-questions for genuine gaps in existing knowledge.
+- The question must be standalone and self-contained
-- sub_questions must be a list of plain strings (max 3).
+- Avoid repeating questions that have already been answered
-- Each sub_question must be standalone and self-contained.
+- Include concrete entities, scope, and any qualifiers
-- Prioritize the highest-value gaps first."""
+
+Provide brief reasoning explaining your decision."""
 
 SEARCH_PROMPT = """You are a search and question-answering specialist.
 
@@ -87,27 +87,6 @@ Guidelines:
 - Be concise and direct; avoid meta commentary about the process.
 - Results are ordered by relevance, with rank 1 being most relevant."""
 
-DECISION_PROMPT = """You are the research evaluator responsible for assessing
-whether gathered evidence sufficiently answers the research question.
-
-Inputs available:
-- Original research question
-- Question-answer pairs with supporting sources
-- Previous evaluation (if any)
-
-Tasks:
-1. Assess whether the collected evidence answers the original question.
-2. Provide a confidence_score in [0,1] reflecting coverage and evidence quality.
-3. Optionally propose up to 3 new sub-questions if important gaps remain.
-
-Output fields:
-- is_sufficient: true when the question is adequately answered
-- confidence_score: numeric in [0,1]
-- reasoning: brief explanation of the assessment
-- new_questions: list of follow-up questions (max 3), only if needed
-
-Be strict: only mark sufficient when key aspects are addressed with reliable evidence."""
-
 SYNTHESIS_PROMPT = """You are a synthesis specialist producing the final
 research report that directly answers the original question.
 

haiku_rag_slim/haiku/rag/agents/research/state.py

@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING
 from pydantic import BaseModel, Field
 
 from haiku.rag.agents.research.dependencies import ResearchContext
-from haiku.rag.agents.research.models import EvaluationResult
 from haiku.rag.client import HaikuRAG
 
 if TYPE_CHECKING:
@@ -30,15 +29,9 @@ class ResearchState(BaseModel):
     )
     iterations: int = Field(default=0, description="Current iteration number")
     max_iterations: int = Field(default=3, description="Maximum allowed iterations")
-    confidence_threshold: float = Field(
-        default=0.8, description="Confidence threshold for completion", ge=0.0, le=1.0
-    )
     max_concurrency: int = Field(
         default=1, description="Maximum concurrent search operations", ge=1
     )
-    last_eval: EvaluationResult | None = Field(
-        default=None, description="Last evaluation result"
-    )
     search_filter: str | None = Field(
         default=None, description="SQL WHERE clause to filter search results"
     )
@@ -49,7 +42,6 @@ class ResearchState(BaseModel):
         context: ResearchContext,
         config: "AppConfig",
         max_iterations: int | None = None,
-        confidence_threshold: float | None = None,
     ) -> "ResearchState":
         """Create a ResearchState from an AppConfig.
 
@@ -57,15 +49,11 @@ class ResearchState(BaseModel):
             context: The ResearchContext containing the question
             config: The AppConfig object
             max_iterations: Override max iterations (None uses config default)
-            confidence_threshold: Override threshold (None uses config, 0.0 disables check)
         """
         return cls(
             context=context,
             max_iterations=max_iterations
             if max_iterations is not None
             else config.research.max_iterations,
-            confidence_threshold=confidence_threshold
-            if confidence_threshold is not None
-            else config.research.confidence_threshold,
             max_concurrency=config.research.max_concurrency,
         )

haiku_rag_slim/haiku/rag/app.py

@@ -398,8 +398,7 @@ class HaikuRAGApp:
                 state = ResearchState.from_config(
                     context=context,
                     config=self.config,
-                    max_iterations=2,
+                    max_iterations=1,
-                    confidence_threshold=0.0,
                 )
                 state.search_filter = filter
                 deps = ResearchDeps(client=self.client)
@@ -477,12 +476,6 @@ 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
-                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]")

haiku_rag_slim/haiku/rag/config/models.py

@@ -78,7 +78,6 @@ class QAConfig(BaseModel):
             enable_thinking=False,
         )
     )
-    max_sub_questions: int = 3
     max_iterations: int = 2
     max_concurrency: int = 1
 
@@ -92,7 +91,6 @@ class ResearchConfig(BaseModel):
         )
     )
     max_iterations: int = 3
-    confidence_threshold: float = 0.8
     max_concurrency: int = 1
 
 

haiku_rag_slim/haiku/rag/mcp.py

@@ -199,7 +199,6 @@ def create_mcp_server(
                         context=context,
                         config=config,
                         max_iterations=2,
-                        confidence_threshold=0.0,
                     )
                     deps = ResearchDeps(client=rag)
 

tests/agents/research/test_models.py

@@ -120,62 +120,3 @@ class TestSearchAnswerPrimarySource:
             citations=[],
         )
         assert answer.primary_source is None
-
-
-class TestFormatContextMerged:
-    """Tests for merged format_context_for_prompt function."""
-
-    def test_format_context_includes_pending_questions_by_default(self):
-        """Test format_context_for_prompt includes pending_questions by default."""
-        from haiku.rag.agents.research.dependencies import ResearchContext
-        from haiku.rag.agents.research.graph import format_context_for_prompt
-
-        context = ResearchContext(
-            original_question="What is X?",
-            sub_questions=["What is A?", "What is B?"],
-        )
-        result = format_context_for_prompt(context)
-        assert "<pending_questions>" in result
-        assert "What is A?" in result
-        assert "What is B?" in result
-
-    def test_format_context_excludes_pending_questions_when_flag_false(self):
-        """Test format_context_for_prompt excludes pending_questions when flag is False."""
-        from haiku.rag.agents.research.dependencies import ResearchContext
-        from haiku.rag.agents.research.graph import format_context_for_prompt
-
-        context = ResearchContext(
-            original_question="What is X?",
-            sub_questions=["What is A?", "What is B?"],
-        )
-        result = format_context_for_prompt(context, include_pending_questions=False)
-        assert "<pending_questions>" not in result
-        assert "What is A?" not in result
-
-    def test_format_context_uses_primary_source_helper(self):
-        """Test format_context_for_prompt uses primary_source from SearchAnswer."""
-        from haiku.rag.agents.research.dependencies import ResearchContext
-        from haiku.rag.agents.research.graph import format_context_for_prompt
-
-        context = ResearchContext(
-            original_question="What is X?",
-        )
-        # Add a QA response with citation
-        answer = SearchAnswer(
-            query="What is A?",
-            answer="A is...",
-            confidence=0.9,
-            citations=[
-                Citation(
-                    document_id="doc-1",
-                    chunk_id="chunk-1",
-                    document_uri="test.md",
-                    document_title="Test Document",
-                    content="content",
-                ),
-            ],
-        )
-        context.add_qa_response(answer)
-
-        result = format_context_for_prompt(context)
-        assert "Test Document" in result

tests/agents/research/test_plan_prompt_selection.py

@@ -1,28 +1,31 @@
-from haiku.rag.agents.research.prompts import PLAN_PROMPT, PLAN_PROMPT_WITH_CONTEXT
+from haiku.rag.agents.research.prompts import (
+    ITERATIVE_PLAN_PROMPT,
+    ITERATIVE_PLAN_PROMPT_WITH_CONTEXT,
+)
 
 
-def test_plan_prompt_with_context_does_not_instruct_gather_context():
+def test_iterative_plan_prompt_proposes_first_question():
-    """PLAN_PROMPT_WITH_CONTEXT should not instruct to use gather_context.
+    """ITERATIVE_PLAN_PROMPT should instruct to propose the first question."""
-
+    assert "first question" in ITERATIVE_PLAN_PROMPT.lower()
-    When session context already exists, we don't need to gather context again.
+    assert "is_complete=False" in ITERATIVE_PLAN_PROMPT
-    """
-    assert "gather_context" not in PLAN_PROMPT_WITH_CONTEXT
 
 
-def test_plan_prompt_instructs_gather_context():
+def test_iterative_plan_prompt_with_context_evaluates_evidence():
-    """PLAN_PROMPT should instruct to use gather_context for initial planning."""
+    """ITERATIVE_PLAN_PROMPT_WITH_CONTEXT should evaluate prior answers."""
-    assert "gather_context" in PLAN_PROMPT
+    assert "prior_answers" in ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
+    assert (
+        "evaluat" in ITERATIVE_PLAN_PROMPT_WITH_CONTEXT.lower()
+    )  # matches evaluate/evaluating
 
 
-def test_prompt_selection_uses_context_prompt_with_session_context():
+def test_prompt_selection_uses_context_prompt_with_prior_answers():
-    """When session_context exists, should use PLAN_PROMPT_WITH_CONTEXT."""
+    """When prior_answers exist, should use ITERATIVE_PLAN_PROMPT_WITH_CONTEXT."""
-    has_prior_answers = False
+    has_prior_answers = True
-    has_session_context = True
 
     effective_plan_prompt = (
-        PLAN_PROMPT_WITH_CONTEXT
+        ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
-        if has_prior_answers or has_session_context
+        if has_prior_answers
-        else PLAN_PROMPT
+        else ITERATIVE_PLAN_PROMPT
     )
 
-    assert effective_plan_prompt == PLAN_PROMPT_WITH_CONTEXT
+    assert effective_plan_prompt == ITERATIVE_PLAN_PROMPT_WITH_CONTEXT

tests/agents/research/test_research_graph.py

@@ -30,7 +30,6 @@ async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus):
     state = ResearchState(
         context=ResearchContext(original_question=doc["question"]),
         max_iterations=1,
-        confidence_threshold=0.5,
         max_concurrency=1,
     )
 
@@ -46,22 +45,27 @@ async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus):
     client.close()
 
 
-def test_research_plan_allows_empty_sub_questions():
+def test_iterative_plan_result_model():
-    """Test ResearchPlan accepts empty sub_questions when context is sufficient."""
+    """Test IterativePlanResult model validation."""
-    from haiku.rag.agents.research.models import ResearchPlan
+    from haiku.rag.agents.research.models import IterativePlanResult
 
-    plan = ResearchPlan(sub_questions=[])
+    # Test complete state
-    assert plan.sub_questions == []
+    complete = IterativePlanResult(
+        is_complete=True,
+        next_question=None,
+        reasoning="All aspects covered.",
+    )
+    assert complete.is_complete is True
+    assert complete.next_question is None
 
-
+    # Test continue state
-def test_research_plan_rejects_too_many_sub_questions():
+    continue_result = IterativePlanResult(
-    """Test ResearchPlan rejects more than 12 sub_questions."""
+        is_complete=False,
-    from pydantic import ValidationError
+        next_question="What are the specific requirements?",
-
+        reasoning="Need more details.",
-    from haiku.rag.agents.research.models import ResearchPlan
+    )
-
+    assert continue_result.is_complete is False
-    with pytest.raises(ValidationError, match="Cannot have more than 12"):
+    assert continue_result.next_question == "What are the specific requirements?"
-        ResearchPlan(sub_questions=[f"q{i}" for i in range(13)])
 
 
 # =============================================================================
@@ -69,13 +73,20 @@ def test_research_plan_rejects_too_many_sub_questions():
 # =============================================================================
 
 
-def test_build_conversational_graph_returns_graph():
+def test_build_research_graph_conversational_mode_returns_graph():
-    """Test build_conversational_graph returns a valid Graph instance."""
+    """Test build_research_graph with output_mode='conversational' returns a valid Graph instance."""
     from pydantic_graph.beta import Graph
 
-    from haiku.rag.agents.research.graph import build_conversational_graph
+    graph = build_research_graph(output_mode="conversational")
+    assert graph is not None
+    assert isinstance(graph, Graph)
 
-    graph = build_conversational_graph()
+
+def test_build_research_graph_report_mode_returns_graph():
+    """Test build_research_graph with output_mode='report' returns a valid Graph instance."""
+    from pydantic_graph.beta import Graph
+
+    graph = build_research_graph(output_mode="report")
     assert graph is not None
     assert isinstance(graph, Graph)
 
@@ -143,27 +154,6 @@ def test_format_context_for_prompt_with_session_context():
     assert "What is Y?" in result
 
 
-def test_format_context_for_prompt_excludes_pending_questions():
-    """Test format_context_for_prompt can exclude pending questions."""
-    from haiku.rag.agents.research.dependencies import ResearchContext
-    from haiku.rag.agents.research.graph import format_context_for_prompt
-
-    context = ResearchContext(
-        original_question="Main question?",
-        sub_questions=["Sub Q1?", "Sub Q2?"],
-    )
-
-    # With pending questions (default)
-    with_pending = format_context_for_prompt(context, include_pending_questions=True)
-    assert "Sub Q1?" in with_pending
-
-    # Without pending questions (for synthesis)
-    without_pending = format_context_for_prompt(
-        context, include_pending_questions=False
-    )
-    assert "Sub Q1?" not in without_pending
-
-
 def test_format_context_for_prompt_with_prior_answers():
     """Test format_context_for_prompt includes prior_answers."""
     from haiku.rag.agents.research.dependencies import ResearchContext

tests/agents/research/test_search_filter.py

@@ -78,7 +78,6 @@ async def test_research_graph_uses_search_filter(
     state = ResearchState(
         context=ResearchContext(original_question="Tell me about animals"),
         max_iterations=1,
-        confidence_threshold=0.5,
         search_filter=filter_clause,
     )
 
@@ -116,7 +115,6 @@ async def test_search_filter_none_searches_all(allow_model_requests, client_with
     state = ResearchState(
         context=ResearchContext(original_question="Tell me about animals"),
         max_iterations=1,
-        confidence_threshold=0.5,
         search_filter=None,
     )
 

tests/cassettes/test_chat_agent/test_chat_agent_get_document_not_found.yaml

@@ -8,7 +8,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '4073'
+      - '5211'
       content-type:
       - application/json
       host:
@@ -23,14 +23,16 @@ interactions:
 
           CRITICAL RULES:
           1. For greetings or casual chat: respond directly WITHOUT using any tools
-          2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
+          2. For questions: Use the "ask" tool EXACTLY ONCE - it automatically uses prior conversation context
           3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
           4. NEVER call the same tool multiple times for a single user message
           5. NEVER make up information - always use tools to get facts from the knowledge base
 
           How to decide which tool to use:
-          - "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
+          - "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
-          - "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
+          - "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
+          - "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document").
+          - "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return answers with citations.
           - "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
 
           IMPORTANT - When user mentions a document in search/ask:
@@ -107,173 +109,17 @@ interactions:
         type: function
       - function:
           description: |-
-            Retrieve a specific document by title or URI.
+            List available documents in the knowledge base.
 
-            Use this when the user wants to fetch/get/retrieve a specific document.
+            Use this when the user wants to browse or see what documents are available.
-          name: get_document
+          name: list_documents
           parameters:
             additionalProperties: false
             properties:
-              query:
+              page:
-                description: The document title or URI to look up
+                default: 1
-                type: string
+                description: 'Page number (default: 1, 50 documents per page)'
-            required:
+                type: integer
-            - query
-            type: object
-          strict: true
-        type: function
-    uri: http://localhost:11434/v1/chat/completions
-  response:
-    headers:
-      content-length:
-      - '559'
-      content-type:
-      - application/json
-    parsed_body:
-      choices:
-      - finish_reason: tool_calls
-        index: 0
-        message:
-          content: ''
-          reasoning: User wants nonexistent document. We can use get_document but it may not exist. We'll try get_document.
-          role: assistant
-          tool_calls:
-          - function:
-              arguments: '{"query":"nonexistent document"}'
-              name: get_document
-            id: call_31uy8050
-            index: 0
-            type: function
-      created: 1768998264
-      id: chatcmpl-114
-      model: gpt-oss
-      object: chat.completion
-      system_fingerprint: fp_ollama
-      usage:
-        completion_tokens: 47
-        prompt_tokens: 842
-        total_tokens: 889
-    status:
-      code: 200
-      message: OK
-- request:
-    headers:
-      accept:
-      - application/json
-      accept-encoding:
-      - gzip, deflate, zstd
-      connection:
-      - keep-alive
-      content-length:
-      - '4470'
-      content-type:
-      - application/json
-      host:
-      - localhost:11434
-    method: POST
-    parsed_body:
-      messages:
-      - content: |-
-          You are a helpful research assistant powered by haiku.rag, a knowledge base system.
-
-          You have access to a knowledge base of documents. Use your tools to search and answer questions.
-
-          CRITICAL RULES:
-          1. For greetings or casual chat: respond directly WITHOUT using any tools
-          2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
-          3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
-          4. NEVER call the same tool multiple times for a single user message
-          5. NEVER make up information - always use tools to get facts from the knowledge base
-
-          How to decide which tool to use:
-          - "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
-          - "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
-          - "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
-
-          IMPORTANT - When user mentions a document in search/ask:
-          - If user says "search in <doc>", "find in <doc>", "answer from <doc>", or "<topic> in <doc>":
-            - Extract the TOPIC as `query`/`question`
-            - Extract the DOCUMENT NAME as `document_name`
-          - Examples for search:
-            - "search for embeddings in the ML paper" → query="embeddings", document_name="ML paper"
-            - "find transformer architecture in 2412.00566" → query="transformer architecture", document_name="2412.00566"
-          - Examples for ask:
-            - "what does the ML paper say about embeddings?" → question="what are the embedding methods?", document_name="ML paper"
-            - "answer from 2412.00566 about model training" → question="how is the model trained?", document_name="2412.00566"
-
-          Be friendly and conversational. When you use the "ask" tool, summarize the key findings for the user.
-        role: system
-      - content: Get me the nonexistent document
-        role: user
-      - content: |-
-          <think>
-          User wants nonexistent document. We can use get_document but it may not exist. We'll try get_document.
-          </think>
-        role: assistant
-        tool_calls:
-        - function:
-            arguments: '{"query":"nonexistent document"}'
-            name: get_document
-          id: call_31uy8050
-          type: function
-      - content: 'Document not found: nonexistent document'
-        role: tool
-        tool_call_id: call_31uy8050
-      model: gpt-oss
-      reasoning_effort: low
-      stream: false
-      tool_choice: auto
-      tools:
-      - function:
-          description: |-
-            Search the knowledge base for relevant documents.
-
-            Use this when you need to find documents or explore the knowledge base.
-            Results are displayed to the user - just list the titles found.
-          name: search
-          parameters:
-            additionalProperties: false
-            properties:
-              document_name:
-                anyOf:
-                - type: string
-                - type: 'null'
-                default: null
-                description: Optional document name/title to search within
-              limit:
-                anyOf:
-                - type: integer
-                - type: 'null'
-                default: null
-                description: 'Number of results to return (default: 5)'
-              query:
-                description: The search query (what to search for)
-                type: string
-            required:
-            - query
-            type: object
-        type: function
-      - function:
-          description: |-
-            Answer a specific question using the knowledge base.
-
-            Use this for direct questions that need a focused answer with citations.
-            Uses a research graph for planning, searching, and synthesis.
-          name: ask
-          parameters:
-            additionalProperties: false
-            properties:
-              document_name:
-                anyOf:
-                - type: string
-                - type: 'null'
-                default: null
-                description: Optional document name/title to search within (e.g., "tbmed593", "army manual")
-              question:
-                description: The question to answer
-                type: string
-            required:
-            - question
             type: object
         type: function
       - function:
@@ -293,11 +139,28 @@ interactions:
             type: object
           strict: true
         type: function
+      - function:
+          description: |-
+            Generate a summary of a specific document.
+
+            Use this when the user wants an overview or summary of a document's content.
+          name: summarize_document
+          parameters:
+            additionalProperties: false
+            properties:
+              query:
+                description: The document title or URI to summarize
+                type: string
+            required:
+            - query
+            type: object
+          strict: true
+        type: function
     uri: http://localhost:11434/v1/chat/completions
   response:
     headers:
       content-length:
-      - '459'
+      - '539'
       content-type:
       - application/json
     parsed_body:
@@ -305,18 +168,19 @@ interactions:
       - finish_reason: stop
         index: 0
         message:
-          content: I’m sorry, but I couldn’t find a document titled “nonexistent document.” If you have another title or some
+          content: I’m sorry, but that document isn’t available in the knowledge base. If there’s another topic or document
-            details to share, let me know and I’ll look it up for you!
+            you’d like help with, just let me know!
+          reasoning: User asking for nonexistent document. Need to respond that none exists. No tool usage.
           role: assistant
-      created: 1768998265
+      created: 1769793913
-      id: chatcmpl-968
+      id: chatcmpl-124
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 44
+        completion_tokens: 60
-        prompt_tokens: 912
+        prompt_tokens: 1025
-        total_tokens: 956
+        total_tokens: 1085
     status:
       code: 200
       message: OK

tests/cassettes/test_chat_agent/test_chat_agent_get_document_tool.yaml

@@ -62,7 +62,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '4084'
+      - '5222'
       content-type:
       - application/json
       host:
@@ -77,14 +77,16 @@ interactions:
 
           CRITICAL RULES:
           1. For greetings or casual chat: respond directly WITHOUT using any tools
-          2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
+          2. For questions: Use the "ask" tool EXACTLY ONCE - it automatically uses prior conversation context
           3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
           4. NEVER call the same tool multiple times for a single user message
           5. NEVER make up information - always use tools to get facts from the knowledge base
 
           How to decide which tool to use:
-          - "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
+          - "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
-          - "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
+          - "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
+          - "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document").
+          - "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return answers with citations.
           - "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
 
           IMPORTANT - When user mentions a document in search/ask:
@@ -159,6 +161,21 @@ interactions:
             - question
             type: object
         type: function
+      - function:
+          description: |-
+            List available documents in the knowledge base.
+
+            Use this when the user wants to browse or see what documents are available.
+          name: list_documents
+          parameters:
+            additionalProperties: false
+            properties:
+              page:
+                default: 1
+                description: 'Page number (default: 1, 50 documents per page)'
+                type: integer
+            type: object
+        type: function
       - function:
           description: |-
             Retrieve a specific document by title or URI.
@@ -176,11 +193,28 @@ interactions:
             type: object
           strict: true
         type: function
+      - function:
+          description: |-
+            Generate a summary of a specific document.
+
+            Use this when the user wants an overview or summary of a document's content.
+          name: summarize_document
+          parameters:
+            additionalProperties: false
+            properties:
+              query:
+                description: The document title or URI to summarize
+                type: string
+            required:
+            - query
+            type: object
+          strict: true
+        type: function
     uri: http://localhost:11434/v1/chat/completions
   response:
     headers:
       content-length:
-      - '477'
+      - '510'
       content-type:
       - application/json
     parsed_body:
@@ -189,24 +223,24 @@ interactions:
         index: 0
         message:
           content: ''
-          reasoning: Need get_document.
+          reasoning: User wants a specific document. Use get_document.
           role: assistant
           tool_calls:
           - function:
               arguments: '{"query":"DocLayNet Class Labels"}'
               name: get_document
-            id: call_m5izt73j
+            id: call_jxvkrbku
             index: 0
             type: function
-      created: 1768998258
+      created: 1769793902
-      id: chatcmpl-390
+      id: chatcmpl-566
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 32
+        completion_tokens: 38
-        prompt_tokens: 846
+        prompt_tokens: 1029
-        total_tokens: 878
+        total_tokens: 1067
     status:
       code: 200
       message: OK
@@ -219,7 +253,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '5169'
+      - '6338'
       content-type:
       - application/json
       host:
@@ -234,14 +268,16 @@ interactions:
 
           CRITICAL RULES:
           1. For greetings or casual chat: respond directly WITHOUT using any tools
-          2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
+          2. For questions: Use the "ask" tool EXACTLY ONCE - it automatically uses prior conversation context
           3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
           4. NEVER call the same tool multiple times for a single user message
           5. NEVER make up information - always use tools to get facts from the knowledge base
 
           How to decide which tool to use:
-          - "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
+          - "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
-          - "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
+          - "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
+          - "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document").
+          - "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return answers with citations.
           - "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
 
           IMPORTANT - When user mentions a document in search/ask:
@@ -261,21 +297,21 @@ interactions:
         role: user
       - content: |-
           <think>
-          Need get_document.
+          User wants a specific document. Use get_document.
           </think>
         role: assistant
         tool_calls:
         - function:
             arguments: '{"query":"DocLayNet Class Labels"}'
             name: get_document
-          id: call_m5izt73j
+          id: call_jxvkrbku
           type: function
       - content: |-
           **DocLayNet Class Labels**
 
-          - ID: cd0fc5b8-ec16-4f1c-99fe-20979ecaab1c
+          - ID: 99c3503e-a8ad-4116-a8be-63fb8048dceb
           - URI: doclaynet-labels
-          - Created: 2026-01-21 14:24
+          - Created: 2026-01-30 19:25
 
           **Content:**
           DocLayNet Dataset - Class Labels
@@ -296,7 +332,7 @@ interactions:
 
           The Text class has the highest count with 510,377 instances in the dataset.
         role: tool
-        tool_call_id: call_m5izt73j
+        tool_call_id: call_jxvkrbku
       model: gpt-oss
       reasoning_effort: low
       stream: false
@@ -354,6 +390,21 @@ interactions:
             - question
             type: object
         type: function
+      - function:
+          description: |-
+            List available documents in the knowledge base.
+
+            Use this when the user wants to browse or see what documents are available.
+          name: list_documents
+          parameters:
+            additionalProperties: false
+            properties:
+              page:
+                default: 1
+                description: 'Page number (default: 1, 50 documents per page)'
+                type: integer
+            type: object
+        type: function
       - function:
           description: |-
             Retrieve a specific document by title or URI.
@@ -371,11 +422,28 @@ interactions:
             type: object
           strict: true
         type: function
+      - function:
+          description: |-
+            Generate a summary of a specific document.
+
+            Use this when the user wants an overview or summary of a document's content.
+          name: summarize_document
+          parameters:
+            additionalProperties: false
+            properties:
+              query:
+                description: The document title or URI to summarize
+                type: string
+            required:
+            - query
+            type: object
+          strict: true
+        type: function
     uri: http://localhost:11434/v1/chat/completions
   response:
     headers:
       content-length:
-      - '1131'
+      - '1137'
       content-type:
       - application/json
     parsed_body:
@@ -383,24 +451,24 @@ interactions:
       - finish_reason: stop
         index: 0
         message:
-          content: "Here’s the full **DocLayNet Class Labels** document:\n\n---\n\n## DocLayNet Dataset – Class Labels\n\nDocLayNet
+          content: "Here’s the **DocLayNet Class Labels** document:\n\n**DocLayNet Dataset – Class Labels**\n\nDocLayNet defines
-            defines **11 distinct class labels** for document layout analysis:\n\n1. **Caption** – Text describing figures
+            **11** distinct class labels for document layout analysis:\n\n1. **Caption** – Text describing figures or tables
-            or tables  \n2. **Footnote** – Notes at the bottom of pages  \n3. **Formula** – Mathematical expressions  \n4.
+            \ \n2. **Footnote** – Notes at the bottom of pages  \n3. **Formula** – Mathematical expressions  \n4. **List‑item**
-            **List‑item** – Items in bulleted or numbered lists  \n5. **Page‑footer** – Footer content on pages  \n6. **Page‑header**
+            – Items in bulleted or numbered lists  \n5. **Page‑footer** – Footer content on pages  \n6. **Page‑header** –
-            – Header content on pages  \n7. **Picture** – Images and diagrams  \n8. **Section‑header** – Headings for document
+            Header content on pages  \n7. **Picture** – Images and diagrams  \n8. **Section‑header** – Headings for document
             sections  \n9. **Table** – Tabular data  \n10. **Text** – Regular paragraph text (highest count: **510,377** instances)
-            \ \n11. **Title** – Document titles  \n\nThe **Text** class has the highest count with 510,377 instances in the
+            \ \n11. **Title** – Document titles  \n\nThe **Text** class has the largest number of instances, with **510,377**
-            dataset."
+            entries in the dataset."
           role: assistant
-      created: 1768998262
+      created: 1769793910
-      id: chatcmpl-746
+      id: chatcmpl-948
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 202
+        completion_tokens: 204
-        prompt_tokens: 1110
+        prompt_tokens: 1297
-        total_tokens: 1312
+        total_tokens: 1501
     status:
       code: 200
       message: OK