From 714c6c627e21023944ba74e7cda0f15d32d8d405 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 17 Sep 2025 17:08:15 +0300 Subject: [PATCH] Presearch Survey agent --- docs/agents.md | 2 ++ src/haiku/rag/research/__init__.py | 2 ++ src/haiku/rag/research/base.py | 10 +++++- src/haiku/rag/research/orchestrator.py | 39 +++++++++++++++++++++-- src/haiku/rag/research/presearch_agent.py | 34 ++++++++++++++++++++ src/haiku/rag/research/prompts.py | 13 ++++++++ 6 files changed, 97 insertions(+), 3 deletions(-) create mode 100644 src/haiku/rag/research/presearch_agent.py diff --git a/docs/agents.md b/docs/agents.md index 41a52c65..33579666 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -43,6 +43,8 @@ The research workflow coordinates specialized agents to plan, search, analyze, a Components: - Orchestrator: Plans, coordinates, and loops until confidence is sufficient +- Presearch Survey: Runs a quick KB scan and summarizes relevant chunk text to + ground the initial plan (plain-text summary; no URIs or scores) - Search Specialist: Performs targeted RAG searches and answers sub‑questions - Analysis & Evaluation: Extracts insights, identifies gaps, proposes new questions - Synthesis: Produces a final structured research report diff --git a/src/haiku/rag/research/__init__.py b/src/haiku/rag/research/__init__.py index e0716ee8..e1e4f1e4 100644 --- a/src/haiku/rag/research/__init__.py +++ b/src/haiku/rag/research/__init__.py @@ -12,6 +12,7 @@ from haiku.rag.research.evaluation_agent import ( EvaluationResult, ) from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan +from haiku.rag.research.presearch_agent import PresearchSurveyAgent from haiku.rag.research.search_agent import SearchSpecialistAgent from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent @@ -25,6 +26,7 @@ __all__ = [ # Specialized agents "SearchAnswer", "SearchSpecialistAgent", + "PresearchSurveyAgent", "AnalysisEvaluationAgent", "EvaluationResult", "SynthesisAgent", diff --git a/src/haiku/rag/research/base.py b/src/haiku/rag/research/base.py index e61b6ab1..1a0796c9 100644 --- a/src/haiku/rag/research/base.py +++ b/src/haiku/rag/research/base.py @@ -33,10 +33,18 @@ class BaseResearchAgent[T](ABC): # Import deps type lazily to avoid circular import during module load from haiku.rag.research.dependencies import ResearchDependencies + # If the agent is expected to return plain text, pass `str` directly. + # Otherwise, wrap the model with ToolOutput for robust tool-handling retries. + agent_output_type: Any + if self.output_type is str: # plain text output + agent_output_type = str + else: + agent_output_type = ToolOutput(self.output_type, max_retries=3) + self._agent = Agent( model=model_obj, deps_type=ResearchDependencies, - output_type=ToolOutput(self.output_type, max_retries=3), + output_type=agent_output_type, system_prompt=self.get_system_prompt(), ) diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index ffde5736..666ed3cf 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -12,6 +12,7 @@ from haiku.rag.research.evaluation_agent import ( AnalysisEvaluationAgent, EvaluationResult, ) +from haiku.rag.research.presearch_agent import PresearchSurveyAgent from haiku.rag.research.prompts import ORCHESTRATOR_PROMPT from haiku.rag.research.search_agent import SearchSpecialistAgent from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent @@ -41,6 +42,9 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): self.search_agent: SearchSpecialistAgent = SearchSpecialistAgent( provider, model ) + self.presearch_agent: PresearchSurveyAgent = PresearchSurveyAgent( + provider, model + ) self.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent( provider, model ) @@ -61,7 +65,12 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): "original_question": context.original_question, "unanswered_questions": context.sub_questions, "qa_responses": [ - {"question": qa.query, "answer": qa.answer} + { + "question": qa.query, + "answer": qa.answer, + "context_snippets": qa.context, + "sources": qa.sources, + } for qa in context.qa_responses ], "insights": context.insights, @@ -99,12 +108,38 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): # Use provided console or create a new one console = console or Console() if verbose else None + # Run a simple presearch survey to summarize KB context + if console: + console.print( + "\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]" + ) + + presearch_result = await self.presearch_agent.run(question, deps=deps) + # Create initial research plan if console: console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]") + # Include the presearch summary to ground the planning step. + + planning_context_xml = format_as_xml( + { + "original_question": question, + "presearch_summary": presearch_result.output or "", + }, + root_tag="planning_context", + ) + + plan_prompt = ( + "Create a research plan for the main question below.\n\n" + f"Main question: {question}\n\n" + "Use this brief presearch summary to inform the plan. Focus the 3 sub-questions " + "on the most important aspects not already obvious from the current KB context.\n\n" + f"{planning_context_xml}" + ) + plan_result: AgentRunResult[ResearchPlan] = await self.run( - f"Create a research plan for: {question}", deps=deps + plan_prompt, deps=deps ) context.sub_questions = plan_result.output.sub_questions diff --git a/src/haiku/rag/research/presearch_agent.py b/src/haiku/rag/research/presearch_agent.py new file mode 100644 index 00000000..8c6e7c33 --- /dev/null +++ b/src/haiku/rag/research/presearch_agent.py @@ -0,0 +1,34 @@ +from pydantic_ai import RunContext +from pydantic_ai.run import AgentRunResult + +from haiku.rag.research.base import BaseResearchAgent +from haiku.rag.research.dependencies import ResearchDependencies +from haiku.rag.research.prompts import PRESEARCH_AGENT_PROMPT + + +class PresearchSurveyAgent(BaseResearchAgent[str]): + """Presearch agent that gathers verbatim context and summarizes it.""" + + def __init__(self, provider: str, model: str) -> None: + super().__init__(provider, model, str) + + async def run( + self, prompt: str, deps: ResearchDependencies, **kwargs + ) -> AgentRunResult[str]: + return await super().run(prompt, deps, **kwargs) + + def get_system_prompt(self) -> str: + return PRESEARCH_AGENT_PROMPT + + def register_tools(self) -> None: + @self.agent.tool + async def gather_context( + ctx: RunContext[ResearchDependencies], + query: str, + limit: int = 6, + ) -> str: + """Return verbatim concatenation of relevant chunk texts.""" + query = query.replace('"', "") + results = await ctx.deps.client.search(query, limit=limit) + expanded = await ctx.deps.client.expand_context(results) + return "\n\n".join(chunk.content for chunk, _ in expanded) diff --git a/src/haiku/rag/research/prompts.py b/src/haiku/rag/research/prompts.py index 4bcec6a6..af5038f7 100644 --- a/src/haiku/rag/research/prompts.py +++ b/src/haiku/rag/research/prompts.py @@ -114,3 +114,16 @@ Focus on creating a report that provides clear value to the reader by: - Highlighting the most important findings - Explaining the implications of the research - Suggesting concrete next steps""" + +PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor. + +Task: +- Call the gather_context tool once with the main question to obtain a + relevant texts from the Knowledge Base (KB). +- Read that context and produce a brief natural-language summary describing + what the KB appears to contain relative to the question. + +Rules: +- Base the summary strictly on the provided text; do not invent. +- Output only the summary as plain text (one short paragraph). +"""