Simplify research planning: remove gather_context, always use planner

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

@@ -13,7 +13,6 @@ from haiku.rag.agents.research.models import (
     RawSearchAnswer,
     ResearchReport,
     SearchAnswer,
-    resolve_citations,
 )
 from haiku.rag.agents.research.prompts import (
     CONVERSATIONAL_SYNTHESIS_PROMPT,
@@ -69,10 +68,17 @@ async def _iterative_plan_logic(
     config: AppConfig,
 ) -> IterativePlanResult:
     """Evaluate context and decide next question or mark complete."""
-    model_config = config.research.model
-
     has_prior_answers = bool(state.context.qa_responses)
-    has_session_context = bool(state.context.session_context)
+
+    # 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
 
     if has_prior_answers:
         effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT_WITH_CONTEXT, config)
@@ -88,39 +94,6 @@ async def _iterative_plan_logic(
         deps_type=ResearchDependencies,
     )
 
-    search_filter = state.search_filter
-
-    # Register gather_context tool only on first iteration (no prior answers)
-    if not has_prior_answers:
-
-        @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)
-            content = "\n\n".join(r.content for r in results)
-
-            # 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)
@@ -128,18 +101,23 @@ async def _iterative_plan_logic(
             f"Review the gathered evidence and decide whether to continue or synthesize.\n\n"
             f"{context_xml}"
         )
-    elif has_session_context:
-        context_xml = format_context_for_prompt(state.context)
-        prompt = f"Explore the knowledge base and plan research.\n\n{context_xml}"
     else:
-        prompt = (
-            f"Explore the knowledge base and plan research.\n\n"
-            f"Main question: {state.context.original_question}"
-        )
+        context_xml = format_context_for_prompt(state.context)
+        prompt = f"Plan the research investigation.\n\n{context_xml}"
 
     agent_deps = ResearchDependencies(client=deps.client, context=state.context)
     result = await plan_agent.run(prompt, deps=agent_deps)
 
+    # 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
 
 
@@ -371,12 +349,7 @@ def build_research_graph(
             .branch(
                 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
-                    ),
+                    matches=lambda r: not r.is_complete and r.next_question is not None,
                 )
                 .label("Continue research")
                 .transform(extract_question)

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

@@ -1,22 +1,24 @@
-ITERATIVE_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.
 
 Your task:
-1. Use the gather_context tool ONCE to explore the knowledge base for information related to the question
-2. Analyze what you find and decide whether to continue or synthesize
+1. Analyze the original question
+2. Propose the first question to investigate
 
-Decision criteria:
-- Set is_complete=True if the gathered context provides sufficient information to answer the question
-- Set is_complete=False with a next_question if you need to investigate a specific aspect further
+For simple questions, investigate them directly. For composite or complex questions,
+you may decompose into a focused sub-question. For example:
+- "What are the benefits and drawbacks of X?" → Start with "What are the benefits of X?"
+- Ambiguous references should be resolved using background context if available
 
-If not complete, propose exactly ONE high-value follow-up question in next_question:
-- The question must be standalone and self-contained
+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
+
+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
-
-Provide brief reasoning explaining your decision."""
+- Avoid ambiguous pronouns (it/they/this/that)"""
 
 ITERATIVE_PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator evaluating gathered evidence.
 

haiku_rag_slim/haiku/rag/app.py

@@ -398,7 +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

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

@@ -4,17 +4,18 @@ from haiku.rag.agents.research.prompts import (
 )
 
 
-def test_iterative_plan_prompt_with_context_does_not_instruct_gather_context():
-    """ITERATIVE_PLAN_PROMPT_WITH_CONTEXT should not instruct to use gather_context.
-
-    When prior answers already exist, we don't need to gather context again.
-    """
-    assert "gather_context" not in ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
+def test_iterative_plan_prompt_proposes_first_question():
+    """ITERATIVE_PLAN_PROMPT should instruct to propose the first question."""
+    assert "first question" in ITERATIVE_PLAN_PROMPT.lower()
+    assert "is_complete=False" in ITERATIVE_PLAN_PROMPT
 
 
-def test_iterative_plan_prompt_instructs_gather_context():
-    """ITERATIVE_PLAN_PROMPT should instruct to use gather_context for initial planning."""
-    assert "gather_context" in ITERATIVE_PLAN_PROMPT
+def test_iterative_plan_prompt_with_context_evaluates_evidence():
+    """ITERATIVE_PLAN_PROMPT_WITH_CONTEXT should evaluate prior answers."""
+    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_prior_answers():