Optimize context prompts

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

@@ -30,49 +30,51 @@ from haiku.rag.utils import build_prompt, get_model
 
 
 def format_context_for_prompt(context: ResearchContext) -> str:
-    """Format the research context as XML for inclusion in prompts."""
-    context_data: dict[str, object] = {
-        "original_question": context.original_question,
-        "unanswered_questions": context.sub_questions,
-        "qa_responses": [
+    """Format the research context as XML for planning prompts."""
+    context_data: dict[str, object] = {}
+
+    if context.initial_context:
+        context_data["background"] = context.initial_context
+
+    context_data["question"] = context.original_question
+
+    if context.sub_questions:
+        context_data["pending_questions"] = context.sub_questions
+
+    if context.qa_responses:
+        context_data["prior_answers"] = [
             {
                 "question": qa.query,
                 "answer": qa.answer,
                 "confidence": qa.confidence,
-                "sources": [
-                    {
-                        "document_uri": c.document_uri,
-                        "document_title": c.document_title,
-                        "page_numbers": c.page_numbers,
-                        "headings": c.headings,
-                    }
-                    for c in qa.citations
-                ],
+                "source": qa.citations[0].document_title or qa.citations[0].document_uri
+                if qa.citations
+                else None,
             }
             for qa in context.qa_responses
-        ],
-    }
-    if context.initial_context:
-        context_data["initial_context"] = context.initial_context
-    return format_as_xml(context_data, root_tag="research_context")
+        ]
+
+    return format_as_xml(context_data, root_tag="context")
 
 
 def format_conversational_context_for_prompt(context: ResearchContext) -> str:
-    """Format context for conversational mode - excludes unanswered_questions."""
-    context_data: dict[str, object] = {
-        "question": context.original_question,
-    }
+    """Format context for synthesis prompts."""
+    context_data: dict[str, object] = {}
 
     if context.initial_context:
-        context_data["initial_context"] = context.initial_context
+        context_data["background"] = context.initial_context
+
+    context_data["question"] = context.original_question
 
-    # Only include conversation_history if there are qa_responses
     if context.qa_responses:
-        context_data["conversation_history"] = [
+        context_data["prior_answers"] = [
             {
                 "question": qa.query,
                 "answer": qa.answer,
-                "sources": [c.document_title or c.document_uri for c in qa.citations],
+                "confidence": qa.confidence,
+                "source": qa.citations[0].document_title or qa.citations[0].document_uri
+                if qa.citations
+                else None,
             }
             for qa in context.qa_responses
         ]
@@ -95,9 +97,12 @@ async def _plan_step_logic(
     model_config = config.research.model
 
     # Use context-aware prompt if we have existing qa_responses
-    has_context = bool(state.context.qa_responses)
+    has_prior_answers = bool(state.context.qa_responses)
+    has_background = bool(state.context.initial_context)
     effective_plan_prompt = (
-        build_prompt(PLAN_PROMPT_WITH_CONTEXT, config) if has_context else plan_prompt
+        build_prompt(PLAN_PROMPT_WITH_CONTEXT, config)
+        if has_prior_answers
+        else plan_prompt
     )
 
     plan_agent = Agent(
@@ -124,13 +129,20 @@ async def _plan_step_logic(
         return "\n\n".join(r.content for r in results)
 
     # Build prompt with existing context if available
-    if has_context:
+    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"{context_xml}\n\n"
             f"Main question: {state.context.original_question}"
         )
+    elif has_background:
+        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 = (
             "Plan a focused approach for the main question.\n\n"

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

@@ -1,10 +1,11 @@
-PLAN_PROMPT = """You are the research orchestrator for a focused, iterative workflow.
+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:
 1. Understand and decompose the main question
 2. Propose a minimal, high-leverage plan
 3. Coordinate specialized agents to gather evidence
-4. Iterate based on gaps and new findings
 
 Plan requirements:
 - Produce at most 3 sub_questions that together cover the main question.
@@ -22,19 +23,22 @@ Use the gather_context tool once on the main question before planning."""
 
 PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator for a focused workflow.
 
-You have access to PREVIOUS CONVERSATION CONTEXT in the qa_responses section below.
-Review this context first - if it already answers the question, generate minimal
-or no sub-questions. Only create sub-questions to fill gaps in the existing context.
+You have access to context that may include:
+- <background>: Domain context for the conversation
+- <prior_answers>: Previous Q&A pairs with confidence scores
+
+Review this first - if prior answers already answer the question completely,
+you may return an empty sub_questions list. Only create sub-questions to
+fill genuine gaps.
 
 Responsibilities:
-1. Review existing qa_responses to understand what's already known
+1. Review prior_answers to understand what's already known
 2. Identify gaps that need additional research
 3. Propose minimal sub-questions only for missing information
 
 Plan requirements:
-- If existing context fully answers the question, return a SINGLE sub-question
-  to verify or slightly expand the answer.
-- Only create new sub-questions for genuine gaps in the existing knowledge.
+- If prior answers fully answer the question, return an empty sub_questions list.
+- Only create new sub-questions for genuine gaps in existing knowledge.
 - sub_questions must be a list of plain strings (max 3).
 - Each sub_question must be standalone and self-contained.
 - Prioritize the highest-value gaps first.
@@ -145,7 +149,8 @@ Output:
 - confidence: Score from 0.0 to 1.0 indicating answer quality.
 
 Guidelines:
-- Base your answer solely on the collected evidence in qa_responses.
+- Base your answer solely on the evidence provided in the context.
+- If a <background> section is provided, use it to frame your answer appropriately.
 - Be thorough - include all relevant information from the evidence.
 - Use formatting (bullet points, numbered lists) when it improves clarity.
 - Do NOT use meta-commentary like "Based on the research..." or "The evidence shows..."

tests/agents/research/test_research_graph.py

@@ -61,45 +61,45 @@ def test_research_context_initial_context_defaults_to_none():
     assert context.initial_context is None
 
 
-def test_format_context_for_prompt_includes_initial_context():
-    """Test format_context_for_prompt includes initial_context in output."""
+def test_format_context_for_prompt_includes_background():
+    """Test format_context_for_prompt includes background in output."""
     from haiku.rag.agents.research.graph import format_context_for_prompt
 
     context = ResearchContext(
         original_question="What is X?",
-        initial_context="Background: X is a concept in domain Y.",
+        initial_context="X is a concept in domain Y.",
     )
     result = format_context_for_prompt(context)
-    assert "Background: X is a concept in domain Y." in result
-    assert "initial_context" in result
+    assert "X is a concept in domain Y." in result
+    assert "<background>" in result
 
 
-def test_format_context_for_prompt_excludes_initial_context_when_none():
-    """Test format_context_for_prompt excludes initial_context when None."""
+def test_format_context_for_prompt_excludes_background_when_none():
+    """Test format_context_for_prompt excludes background when None."""
     from haiku.rag.agents.research.graph import format_context_for_prompt
 
     context = ResearchContext(original_question="What is X?")
     result = format_context_for_prompt(context)
-    assert "initial_context" not in result
+    assert "<background>" not in result
 
 
-def test_format_conversational_context_for_prompt_includes_initial_context():
-    """Test format_conversational_context_for_prompt includes initial_context."""
+def test_format_conversational_context_for_prompt_includes_background():
+    """Test format_conversational_context_for_prompt includes background."""
     from haiku.rag.agents.research.graph import format_conversational_context_for_prompt
 
     context = ResearchContext(
         original_question="What is X?",
-        initial_context="Background: X is a concept in domain Y.",
+        initial_context="X is a concept in domain Y.",
     )
     result = format_conversational_context_for_prompt(context)
-    assert "Background: X is a concept in domain Y." in result
-    assert "initial_context" in result
+    assert "X is a concept in domain Y." in result
+    assert "<background>" in result
 
 
-def test_format_conversational_context_for_prompt_excludes_initial_context_when_none():
-    """Test format_conversational_context_for_prompt excludes initial_context when None."""
+def test_format_conversational_context_for_prompt_excludes_background_when_none():
+    """Test format_conversational_context_for_prompt excludes background when None."""
     from haiku.rag.agents.research.graph import format_conversational_context_for_prompt
 
     context = ResearchContext(original_question="What is X?")
     result = format_conversational_context_for_prompt(context)
-    assert "initial_context" not in result
+    assert "<background>" not in result