Refactor research graph to iterative planning approach

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

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/graph.py

@@ -1,25 +1,24 @@
 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,
+    resolve_citations,
 )
 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,
 )
@@ -64,33 +63,26 @@ 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."""
     model_config = config.research.model
 
-    # Use context-aware prompt if we have existing qa_responses or session_context
     has_prior_answers = bool(state.context.qa_responses)
     has_session_context = bool(state.context.session_context)
-    effective_plan_prompt = (
-        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]
+    if has_prior_answers:
+        effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT_WITH_CONTEXT, config)
+    else:
+        effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT, config)
+
+    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,
@@ -98,8 +90,8 @@ async def _plan_step_logic(
 
     search_filter = state.search_filter
 
-    # Only register gather_context tool when we don't have existing context
+    # Register gather_context tool only on first iteration (no prior answers)
-    if not has_prior_answers and not has_session_context:
+    if not has_prior_answers:
 
         @plan_agent.tool
         async def gather_context(
@@ -111,33 +103,44 @@ async def _plan_step_logic(
                 query, limit=limit, filter=search_filter
             )
             results = await ctx2.deps.client.expand_context(results)
-            return "\n\n".join(r.content for r in results)
+            content = "\n\n".join(r.content for r in results)
 
-    # Build prompt with existing context if available
+            # Save as a preliminary answer so synthesis has context if planner
+            # decides to complete immediately
+            if results:
+                preliminary = SearchAnswer(
+                    query=query,
+                    answer=content,
+                    cited_chunks=[r.chunk_id for r in results if r.chunk_id],
+                    confidence=0.5,
+                    citations=resolve_citations(
+                        [r.chunk_id for r in results if r.chunk_id], results
+                    ),
+                )
+                state.context.add_qa_response(preliminary)
+
+            return content
+
+    # Build prompt based on current state
     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 = (
+        prompt = f"Explore the knowledge base and plan research.\n\n{context_xml}"
-            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 = (
-            "Plan a focused approach for the main question.\n\n"
+            f"Explore the knowledge base and plan research.\n\n"
             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)
+    return result.output
 
 
 async def _search_one_step_logic(
@@ -147,14 +150,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 +179,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 +192,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 +203,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 +270,123 @@ 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, include_pending_questions=False
+                state.context, include_pending_questions=False
-        )
+            )
-        prompt = f"Answer the question based on the gathered evidence.\n\n{context_xml}"
+            prompt = (
-        agent_deps = ResearchDependencies(
+                "Generate a comprehensive research report based on all gathered information.\n\n"
-            client=deps.client,
+                f"{context_xml}\n\n"
-            context=state.context,
+                "Create a detailed report that synthesizes all findings into a coherent response."
-        )
+            )
-        result = await agent.run(prompt, deps=agent_deps)
+            agent_deps = ResearchDependencies(
+                client=deps.client,
+                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, include_pending_questions=False
+            )
+            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, ctx=None: (
+                        not r.is_complete
+                        and r.next_question is not None
+                        and ctx is not None
+                        and ctx.state.iterations < ctx.state.max_iterations
+                    ),
+                )
+                .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,45 @@
-PLAN_PROMPT = """You are the research orchestrator for a focused workflow.
+ITERATIVE_PLAN_PROMPT = """You are the research orchestrator for a focused workflow.
 
 If a <background> section is provided, use it to understand the domain context.
 
-Responsibilities:
+Your task:
-1. Understand and decompose the main question
+1. Use the gather_context tool ONCE to explore the knowledge base for information related to the question
-2. Propose a minimal, high-leverage plan
+2. Analyze what you find and decide whether to continue or synthesize
-3. Coordinate specialized agents to gather evidence
 
-Plan requirements:
+Decision criteria:
-- Produce at most 3 sub_questions that together cover the main question.
+- Set is_complete=True if the gathered context provides sufficient information to answer the question
-- sub_questions must be a list of plain strings, where each string is a complete
+- Set is_complete=False with a next_question if you need to investigate a specific aspect further
-  question. Do NOT use objects with nested fields like {question, details}.
-- 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."""
+If not complete, propose exactly ONE high-value follow-up question in next_question:
+- The question must be standalone and self-contained
+- Include concrete entities, scope, and any qualifiers
+- Avoid ambiguous pronouns (it/they/this/that)
+- Focus on the most important gap in knowledge
 
-PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator for a focused workflow.
+Provide brief reasoning explaining your decision."""
+
+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 +85,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:
@@ -36,9 +35,6 @@ class ResearchState(BaseModel):
     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"
     )

haiku_rag_slim/haiku/rag/app.py

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

tests/agents/research/test_plan_prompt_selection.py

@@ -1,28 +1,30 @@
-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_with_context_does_not_instruct_gather_context():
-    """PLAN_PROMPT_WITH_CONTEXT should not instruct to use gather_context.
+    """ITERATIVE_PLAN_PROMPT_WITH_CONTEXT should not instruct to use gather_context.
 
-    When session context already exists, we don't need to gather context again.
+    When prior answers already exist, we don't need to gather context again.
     """
-    assert "gather_context" not in PLAN_PROMPT_WITH_CONTEXT
+    assert "gather_context" not in ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
 
 
-def test_plan_prompt_instructs_gather_context():
+def test_iterative_plan_prompt_instructs_gather_context():
-    """PLAN_PROMPT should instruct to use gather_context for initial planning."""
+    """ITERATIVE_PLAN_PROMPT should instruct to use gather_context for initial planning."""
-    assert "gather_context" in PLAN_PROMPT
+    assert "gather_context" in ITERATIVE_PLAN_PROMPT
 
 
-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

@@ -46,22 +46,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 +74,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)
 

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