From ff92ee82719629ea7748e24cdad5cde4306feff9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 23 Jan 2026 18:40:35 +0200 Subject: [PATCH] Skip gather_context when session context exists --- .../haiku/rag/agents/research/graph.py | 27 +++++++------- .../haiku/rag/agents/research/prompts.py | 14 ++++---- .../research/test_plan_prompt_selection.py | 35 +++++++++++++++++++ 3 files changed, 56 insertions(+), 20 deletions(-) create mode 100644 tests/agents/research/test_plan_prompt_selection.py diff --git a/haiku_rag_slim/haiku/rag/agents/research/graph.py b/haiku_rag_slim/haiku/rag/agents/research/graph.py index feaaea8d..eefb084f 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/graph.py +++ b/haiku_rag_slim/haiku/rag/agents/research/graph.py @@ -83,7 +83,7 @@ async def _plan_step_logic( has_session_context = bool(state.context.session_context) effective_plan_prompt = ( build_prompt(PLAN_PROMPT_WITH_CONTEXT, config) - if has_prior_answers + if has_prior_answers or has_session_context else plan_prompt ) @@ -98,17 +98,20 @@ async def _plan_step_logic( search_filter = state.search_filter - @plan_agent.tool - async def gather_context( - ctx2: RunContext[ResearchDependencies], - query: str, - limit: int | None = None, - ) -> str: - results = await ctx2.deps.client.search( - query, limit=limit, filter=search_filter - ) - results = await ctx2.deps.client.expand_context(results) - return "\n\n".join(r.content for r in results) + # Only register gather_context tool when we don't have existing context + if not has_prior_answers and not has_session_context: + + @plan_agent.tool + async def gather_context( + ctx2: RunContext[ResearchDependencies], + query: str, + limit: int | None = None, + ) -> str: + results = await ctx2.deps.client.search( + query, limit=limit, filter=search_filter + ) + results = await ctx2.deps.client.expand_context(results) + return "\n\n".join(r.content for r in results) # Build prompt with existing context if available if has_prior_answers: diff --git a/haiku_rag_slim/haiku/rag/agents/research/prompts.py b/haiku_rag_slim/haiku/rag/agents/research/prompts.py index 73bc20fc..287388e7 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/prompts.py +++ b/haiku_rag_slim/haiku/rag/agents/research/prompts.py @@ -27,23 +27,21 @@ 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. +Review the provided context first. Use to understand the domain. +If exist and already answer the question completely, +return an empty sub_questions list. Only create sub-questions to fill gaps. Responsibilities: -1. Review prior_answers to understand what's already known +1. Review provided context 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 prior answers fully answer the question, return an empty sub_questions list. +- If existing context fully answers 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. - -Use the gather_context tool once on the main question before planning.""" +- Prioritize the highest-value gaps first.""" SEARCH_PROMPT = """You are a search and question-answering specialist. diff --git a/tests/agents/research/test_plan_prompt_selection.py b/tests/agents/research/test_plan_prompt_selection.py new file mode 100644 index 00000000..b074a3d2 --- /dev/null +++ b/tests/agents/research/test_plan_prompt_selection.py @@ -0,0 +1,35 @@ +from haiku.rag.agents.research.prompts import PLAN_PROMPT, PLAN_PROMPT_WITH_CONTEXT + + +def test_plan_prompt_with_context_does_not_instruct_gather_context(): + """PLAN_PROMPT_WITH_CONTEXT should not instruct to use gather_context. + + When session context already exists, we don't need to gather context again. + """ + assert "gather_context" not in PLAN_PROMPT_WITH_CONTEXT + + +def test_plan_prompt_instructs_gather_context(): + """PLAN_PROMPT should instruct to use gather_context for initial planning.""" + assert "gather_context" in PLAN_PROMPT + + +def test_prompt_selection_uses_context_prompt_with_session_context(): + """When session_context exists, should use PLAN_PROMPT_WITH_CONTEXT. + + This tests the logic pattern used in _plan_step_logic. + """ + # Simulate the selection logic from graph.py + has_prior_answers = False + has_session_context = True + + # Current buggy behavior would select plan_prompt (the one with gather_context) + # Expected behavior: use PLAN_PROMPT_WITH_CONTEXT when session_context exists + effective_plan_prompt = ( + PLAN_PROMPT_WITH_CONTEXT + if has_prior_answers or has_session_context + else PLAN_PROMPT + ) + + # Since session_context exists, we should use PLAN_PROMPT_WITH_CONTEXT + assert effective_plan_prompt == PLAN_PROMPT_WITH_CONTEXT