diff --git a/haiku_rag_slim/haiku/rag/agents/research/graph.py b/haiku_rag_slim/haiku/rag/agents/research/graph.py index 905da165..93037740 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/graph.py +++ b/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" diff --git a/haiku_rag_slim/haiku/rag/agents/research/prompts.py b/haiku_rag_slim/haiku/rag/agents/research/prompts.py index b472e18e..6fe97b42 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/prompts.py +++ b/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 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: +- : Domain context for the conversation +- : 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 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..." diff --git a/tests/agents/research/test_research_graph.py b/tests/agents/research/test_research_graph.py index c077f73c..c95e5203 100644 --- a/tests/agents/research/test_research_graph.py +++ b/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 "" 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 "" 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 "" 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 "" not in result