Refactor to give its research subagent full responsibility

docs/agents.md

@@ -70,17 +70,14 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.research import ResearchOrchestrator
 
 client = HaikuRAG(path_to_db)
-orchestrator = ResearchOrchestrator(
+orchestrator = ResearchOrchestrator(provider="ollama", model="gpt-oss")
-    provider="ollama",
-    model="gpt-oss",
-    verbose=True
-)
 
 report = await orchestrator.conduct_research(
     question="What are the main drivers and recent trends of global temperature anomalies since 1990?",
     client=client,
     max_iterations=2,
     confidence_threshold=0.8,
+    verbose=True,
 )
 
 print(report.title)

src/haiku/rag/app.py

@@ -86,7 +86,7 @@ class HaikuRAGApp:
         async with HaikuRAG(db_path=self.db_path) as client:
             try:
                 # Create orchestrator with default config or fallback to QA
-                orchestrator = ResearchOrchestrator(verbose=verbose)
+                orchestrator = ResearchOrchestrator()
 
                 if verbose:
                     self.console.print(
@@ -100,6 +100,7 @@ class HaikuRAGApp:
                     question=question,
                     client=client,
                     max_iterations=max_iterations,
+                    verbose=verbose,
                 )
 
                 # Display the report

src/haiku/rag/research/base.py

@@ -45,7 +45,7 @@ class BaseResearchAgent[T](ABC):
             model=model_obj,
             deps_type=ResearchDependencies,
             output_type=agent_output_type,
-            system_prompt=self.get_system_prompt(),
+            instructions=self.get_system_prompt(),
             retries=3,
         )
 

src/haiku/rag/research/dependencies.py

@@ -1,4 +1,6 @@
 from pydantic import BaseModel, Field
+from pydantic_ai import format_as_xml
+from rich.console import Console
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.research.base import SearchAnswer
@@ -43,3 +45,25 @@ class ResearchDependencies(BaseModel):
 
     client: HaikuRAG = Field(description="RAG client for document operations")
     context: ResearchContext = Field(description="Shared research context")
+    console: Console | None = None
+
+
+def _format_context_for_prompt(context: ResearchContext) -> str:
+    """Format the research context as XML for inclusion in prompts."""
+
+    context_data = {
+        "original_question": context.original_question,
+        "unanswered_questions": context.sub_questions,
+        "qa_responses": [
+            {
+                "question": qa.query,
+                "answer": qa.answer,
+                "context_snippets": qa.context,
+                "sources": qa.sources,
+            }
+            for qa in context.qa_responses
+        ],
+        "insights": context.insights,
+        "gaps": context.gaps,
+    }
+    return format_as_xml(context_data, root_tag="research_context")

src/haiku/rag/research/evaluation_agent.py

@@ -1,6 +1,11 @@
 from pydantic import BaseModel, Field
+from pydantic_ai.run import AgentRunResult
 
 from haiku.rag.research.base import BaseResearchAgent
+from haiku.rag.research.dependencies import (
+    ResearchDependencies,
+    _format_context_for_prompt,
+)
 from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
 
 
@@ -34,5 +39,47 @@ class AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]):
     def __init__(self, provider: str, model: str) -> None:
         super().__init__(provider, model, output_type=EvaluationResult)
 
+    async def run(
+        self, prompt: str, deps: ResearchDependencies, **kwargs
+    ) -> AgentRunResult[EvaluationResult]:
+        console = deps.console
+        if console:
+            console.print(
+                "\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]"
+            )
+
+        # Format context for the evaluation agent
+        context_xml = _format_context_for_prompt(deps.context)
+        evaluation_prompt = f"""Analyze all gathered information and evaluate the completeness of research.
+
+{context_xml}
+
+Evaluate the research progress for the original question and identify any remaining gaps."""
+
+        result = await super().run(evaluation_prompt, deps, **kwargs)
+        output = result.output
+
+        # Store insights
+        for insight in output.key_insights:
+            deps.context.add_insight(insight)
+
+        # Add new questions to the sub-questions list
+        for new_q in output.new_questions:
+            if new_q not in deps.context.sub_questions:
+                deps.context.sub_questions.append(new_q)
+
+        if console:
+            if output.key_insights:
+                console.print("   [bold]Key insights:[/bold]")
+                for insight in output.key_insights:
+                    console.print(f"   • {insight}")
+            console.print(
+                f"   Confidence: [yellow]{output.confidence_score:.1%}[/yellow]"
+            )
+            status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
+            console.print(f"   Sufficient: {status}")
+
+        return result
+
     def get_system_prompt(self) -> str:
         return EVALUATION_AGENT_PROMPT

src/haiku/rag/research/orchestrator.py

@@ -1,13 +1,15 @@
 from typing import Any
 
 from pydantic import BaseModel, Field
-from pydantic_ai.format_prompt import format_as_xml
 from pydantic_ai.run import AgentRunResult
 from rich.console import Console
 
 from haiku.rag.config import Config
 from haiku.rag.research.base import BaseResearchAgent
-from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
+from haiku.rag.research.dependencies import (
+    ResearchContext,
+    ResearchDependencies,
+)
 from haiku.rag.research.evaluation_agent import (
     AnalysisEvaluationAgent,
     EvaluationResult,
@@ -34,7 +36,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
         self,
         provider: str | None = Config.RESEARCH_PROVIDER,
         model: str | None = None,
-        verbose: bool = False,
     ):
         # Use provided values or fall back to config defaults
         provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER
@@ -52,34 +53,10 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
             provider, model
         )
         self.synthesis_agent: SynthesisAgent = SynthesisAgent(provider, model)
-        if verbose:
-            self._console = Console()
-
-        self.verbose = verbose
 
     def get_system_prompt(self) -> str:
         return ORCHESTRATOR_PROMPT
 
-    def _format_context_for_prompt(self, context: ResearchContext) -> str:
-        """Format the research context as XML for inclusion in prompts."""
-
-        context_data = {
-            "original_question": context.original_question,
-            "unanswered_questions": context.sub_questions,
-            "qa_responses": [
-                {
-                    "question": qa.query,
-                    "answer": qa.answer,
-                    "context_snippets": qa.context,
-                    "sources": qa.sources,
-                }
-                for qa in context.qa_responses
-            ],
-            "insights": context.insights,
-            "gaps": context.gaps,
-        }
-        return format_as_xml(context_data, root_tag="research_context")
-
     def _should_stop_research(
         self,
         evaluation_result: AgentRunResult[EvaluationResult],
@@ -90,20 +67,13 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
         result = evaluation_result.output
         return result.is_sufficient and result.confidence_score >= confidence_threshold
 
-    def _log(self, line="", rule=False):
-        if not self._console:
-            return
-        if rule:
-            self._console.rule(line)
-        else:
-            self._console.print(line)
-
     async def conduct_research(
         self,
         question: str,
         client: Any,
         max_iterations: int = 3,
         confidence_threshold: float = 0.8,
+        verbose: bool = False,
     ) -> ResearchReport:
         """Conduct comprehensive research on a question.
 
@@ -113,7 +83,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
             max_iterations: Maximum number of search-analyze-clarify cycles
             confidence_threshold: Minimum confidence level to stop research (0-1)
             verbose: If True, print progress and intermediate results
-            console: Optional Rich console for output
 
         Returns:
             ResearchReport with comprehensive findings
@@ -122,14 +91,16 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
         # Initialize context
         context = ResearchContext(original_question=question)
         deps = ResearchDependencies(client=client, context=context)
+        if verbose:
+            deps.console = Console()
 
+        console = deps.console
         # Create initial research plan
+        if console:
+            console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
+
         # Run a simple presearch survey to summarize KB context
-        self._log("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
-        self._log("\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]")
-
         presearch_result = await self.presearch_agent.run(question, deps=deps)
-
         plan_prompt = (
             "Create a research plan for the main question below.\n\n"
             f"Main question: {question}\n\n"
@@ -143,109 +114,57 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
         )
         context.sub_questions = plan_result.output.sub_questions
 
-        self._log("\n[bold green]✅ Research Plan Created:[/bold green]")
+        if console:
-        self._log(f"   [bold]Main Question:[/bold] {plan_result.output.main_question}")
+            console.print("\n[bold green]✅ Research Plan Created:[/bold green]")
-        self._log("   [bold]Sub-questions:[/bold]")
+            console.print(
-        for i, sq in enumerate(plan_result.output.sub_questions, 1):
+                f"   [bold]Main Question:[/bold] {plan_result.output.main_question}"
-            self._log(f"      {i}. {sq}")
+            )
+            console.print("   [bold]Sub-questions:[/bold]")
+            for i, sq in enumerate(plan_result.output.sub_questions, 1):
+                console.print(f"      {i}. {sq}")
 
         # Execute research iterations
         for iteration in range(max_iterations):
-            self._log(
+            if console:
-                f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]",
+                console.rule(
-                rule=True,
+                    f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]"
-            )
+                )
 
             # Check if we have questions to search
             if not context.sub_questions:
-                # No more questions to explore
+                if console:
-                self._log(
+                    console.print(
-                    "[yellow]No more questions to explore. Concluding research.[/yellow]"
+                        "[yellow]No more questions to explore. Concluding research.[/yellow]"
-                )
+                    )
                 break
 
             # Use current sub-questions for this iteration
             questions_to_search = context.sub_questions[:]
 
             # Search phase - answer all questions in this iteration
-            self._log(
+            if console:
-                f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]"
+                console.print(
-            )
+                    f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]"
-            for i, q in enumerate(questions_to_search, 1):
+                )
-                self._log(f"   {i}. {q}")
 
-            # Run searches for all questions and remove answered ones
             for search_question in questions_to_search:
                 await self.search_agent.run(search_question, deps=deps)
-                if self._console and context.qa_responses:
-                    # Show the last QA response (which should be for this question)
-                    latest_qa = context.qa_responses[-1]
-                    answer_preview = (
-                        latest_qa.answer[:150] + "…"
-                        if len(latest_qa.answer) > 150
-                        else latest_qa.answer
-                    )
-                    self._log(f"\n   [green]✓[/green] {search_question}")
-                    self._log(f"      {answer_preview}")
 
             # Analysis and Evaluation phase
-            self._log(
-                "\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]"
-            )
 
-            # Format context for the evaluation agent
+            evaluation_result = await self.evaluation_agent.run("", deps=deps)
-            context_xml = self._format_context_for_prompt(context)
-            evaluation_prompt = f"""Analyze all gathered information and evaluate the completeness of research.
-
-{context_xml}
-
-Evaluate the research progress for the original question and identify any remaining gaps."""
-
-            evaluation_result = await self.evaluation_agent.run(
-                evaluation_prompt,
-                deps=deps,
-            )
-
-            output = evaluation_result.output
-            if output.key_insights:
-                self._log("   [bold]Key insights:[/bold]")
-                for insight in output.key_insights:
-                    self._log(f"   • {insight}")
-            self._log(f"   Confidence: [yellow]{output.confidence_score:.1%}[/yellow]")
-            status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
-            self._log(f"   Sufficient: {status}")
-
-            # Store insights
-            for insight in evaluation_result.output.key_insights:
-                context.add_insight(insight)
-
-            # Add new questions to the sub-questions list
-            for new_q in evaluation_result.output.new_questions:
-                if new_q not in context.sub_questions:
-                    context.sub_questions.append(new_q)
 
             # Check if research is sufficient
             if self._should_stop_research(evaluation_result, confidence_threshold):
-                self._log(
+                if console:
-                    f"\n[bold green]✅ Stopping research:[/bold green] {evaluation_result.output.reasoning}"
+                    console.print(
-                )
+                        f"\n[bold green]✅ Stopping research:[/bold green] {evaluation_result.output.reasoning}"
+                    )
                 break
 
         # Generate final report
-        self._log("\n[bold cyan]📝 Generating final research report...[/bold cyan]")
-
-        # Format context for the synthesis agent
-        final_context_xml = self._format_context_for_prompt(context)
-        synthesis_prompt = f"""Generate a comprehensive research report based on all gathered information.
-
-{final_context_xml}
-
-Create a detailed report that synthesizes all findings into a coherent response."""
-
         report_result: AgentRunResult[ResearchReport] = await self.synthesis_agent.run(
-            synthesis_prompt, deps=deps
+            "", deps=deps
         )
 
-        self._log("[bold green]✅ Research complete![/bold green]")
-
         return report_result.output

src/haiku/rag/research/presearch_agent.py

@@ -15,6 +15,12 @@ class PresearchSurveyAgent(BaseResearchAgent[str]):
     async def run(
         self, prompt: str, deps: ResearchDependencies, **kwargs
     ) -> AgentRunResult[str]:
+        console = deps.console
+        if console:
+            console.print(
+                "\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]"
+            )
+
         return await super().run(prompt, deps, **kwargs)
 
     def get_system_prompt(self) -> str:
@@ -28,7 +34,6 @@ class PresearchSurveyAgent(BaseResearchAgent[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)

src/haiku/rag/research/search_agent.py

@@ -21,9 +21,18 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
         Pydantic AI enforces `SearchAnswer` as the output model; we just store
         the QA response with the last search results as sources.
         """
+        console = deps.console
+        if console:
+            console.print(f"\t{prompt}")
+
         result = await super().run(prompt, deps, **kwargs)
         deps.context.add_qa_response(result.output)
         deps.context.sub_questions.remove(prompt)
+        if console:
+            answer = result.output.answer
+            answer_preview = answer[:150] + "…" if len(answer) > 150 else answer
+            console.log(f"\n   [green]✓[/green] {answer_preview}")
+
         return result
 
     def get_system_prompt(self) -> str:
@@ -39,9 +48,6 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
             limit: int = 5,
         ) -> str:
             """Search the KB and return a concise context pack."""
-            # Remove quotes from queries as this requires positional indexing in lancedb
-            # XXX: Investigate how to do that with lancedb
-            query = query.replace('"', "")
             search_results = await ctx.deps.client.search(query, limit=limit)
             expanded = await ctx.deps.client.expand_context(search_results)
 

src/haiku/rag/research/synthesis_agent.py

@@ -1,6 +1,11 @@
 from pydantic import BaseModel, Field
+from pydantic_ai.run import AgentRunResult
 
 from haiku.rag.research.base import BaseResearchAgent
+from haiku.rag.research.dependencies import (
+    ResearchDependencies,
+    _format_context_for_prompt,
+)
 from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT
 
 
@@ -30,5 +35,26 @@ class SynthesisAgent(BaseResearchAgent[ResearchReport]):
     def __init__(self, provider: str, model: str) -> None:
         super().__init__(provider, model, output_type=ResearchReport)
 
+    async def run(
+        self, prompt: str, deps: ResearchDependencies, **kwargs
+    ) -> AgentRunResult[ResearchReport]:
+        console = deps.console
+        if console:
+            console.print(
+                "\n[bold cyan]📝 Generating final research report...[/bold cyan]"
+            )
+
+        context_xml = _format_context_for_prompt(deps.context)
+        synthesis_prompt = f"""Generate a comprehensive research report based on all gathered information.
+
+{context_xml}
+
+Create a detailed report that synthesizes all findings into a coherent response."""
+        result = await super().run(synthesis_prompt, deps, **kwargs)
+        if console:
+            console.print("[bold green]✅ Research complete![/bold green]")
+
+        return result
+
     def get_system_prompt(self) -> str:
         return SYNTHESIS_AGENT_PROMPT

tests/generate_benchmark_db.py

@@ -12,7 +12,7 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.logging import configure_cli_logging
 from haiku.rag.qa import get_qa_agent
 
-logfire.configure()
+logfire.configure(send_to_logfire="if-token-present")
 logfire.instrument_pydantic_ai()
 configure_cli_logging()
 console = Console()