danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Simplify orchestration & subagents

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-16 14:25:28 +03:00
parent 5327b84fb9
commit 98ca5c4746
No known key found for this signature in database
13 changed files with 374 additions and 561 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
src/haiku/rag/qa/agent.py
View file

@ -6,7 +6,7 @@ from pydantic_ai.providers.openai import OpenAIProvider
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.qa.prompts import SYSTEM_PROMPT, SYSTEM_PROMPT_WITH_CITATIONS
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIONS
class SearchResult(BaseModel):
@ -31,7 +31,9 @@ class QuestionAnswerAgent:
):
self._client = client
system_prompt = SYSTEM_PROMPT_WITH_CITATIONS if use_citations else SYSTEM_PROMPT
system_prompt = (
QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
)
model_obj = self._get_model(provider, model)
self._agent = Agent(

4
src/haiku/rag/qa/prompts.py
View file

@ -1,4 +1,4 @@
SYSTEM_PROMPT = """
QA_SYSTEM_PROMPT = """
You are a knowledgeable assistant that helps users find information from a document knowledge base.
Your process:
@ -21,7 +21,7 @@ Be concise, and always maintain accuracy over completeness. Prefer short, direct
/no_think
"""
SYSTEM_PROMPT_WITH_CITATIONS = """
QA_SYSTEM_PROMPT_WITH_CITATIONS = """
You are a knowledgeable assistant that helps users find information from a document knowledge base.
IMPORTANT: You MUST use the search_documents tool for every question. Do not answer any question without first searching the knowledge base.

15
src/haiku/rag/research/__init__.py
View file

@ -1,12 +1,11 @@
"""Multi-agent research workflow for advanced RAG queries."""
from haiku.rag.research.analysis_agent import AnalysisAgent, AnalysisResult
from haiku.rag.research.base import BaseResearchAgent, ResearchOutput, SearchResult
from haiku.rag.research.clarification_agent import (
ClarificationAgent,
ClarificationResult,
)
from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.research.evaluation_agent import (
AnalysisEvaluationAgent,
EvaluationResult,
)
from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan
from haiku.rag.research.search_agent import SearchSpecialistAgent
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
@ -20,10 +19,8 @@ __all__ = [
"ResearchOutput",
# Specialized agents
"SearchSpecialistAgent",
"AnalysisAgent",
"AnalysisResult",
"ClarificationAgent",
"ClarificationResult",
"AnalysisEvaluationAgent",
"EvaluationResult",
"SynthesisAgent",
"ResearchReport",
# Orchestrator

53
src/haiku/rag/research/analysis_agent.py
View file

@ -1,53 +0,0 @@
"""Analysis agent for content processing and insight extraction."""
from pydantic import BaseModel, Field
from haiku.rag.research.base import BaseResearchAgent
class AnalysisResult(BaseModel):
"""Result of content analysis."""
key_insights: list[str] = Field(
description="Main insights extracted from the documents"
)
themes: dict[str, list[str]] = Field(description="Themes and related findings")
summary: str = Field(description="Consolidated summary of findings")
evidence_quality: str = Field(
description="Assessment of evidence quality (strong/moderate/weak)"
)
recommendations: list[str] = Field(
description="Suggested next steps or areas for further research"
)
class AnalysisAgent(BaseResearchAgent):
"""Agent specialized in content analysis and synthesis."""
def __init__(self, provider: str, model: str):
super().__init__(provider, model, output_type=AnalysisResult)
def get_system_prompt(self) -> str:
return """You are an analysis specialist agent focused on extracting deep insights from search results.
Your role is to:
1. Carefully read and analyze all provided documents
2. Extract key insights and important facts
3. Identify common themes and patterns across documents
4. Synthesize information into a coherent understanding
5. Assess the quality and reliability of the evidence
6. Identify areas that need further investigation
Be specific and detailed in your analysis. Focus on:
- What the documents actually say (not assumptions)
- Connections and contradictions between sources
- The strength of the evidence presented
- Gaps in the information that need to be filled
Your analysis should be thorough, critical, and actionable."""
def register_tools(self) -> None:
"""Register analysis-specific tools."""
# The agent will use its LLM capabilities directly for analysis
# No need for hardcoded tools - the structured output will guide the analysis
pass

17
src/haiku/rag/research/base.py
View file

@ -1,13 +1,12 @@
"""Base class for research agents with common patterns."""
from abc import ABC, abstractmethod
from typing import Any, TypeVar
from typing import Any, Generic, TypeVar
from pydantic import BaseModel
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.ollama import OllamaProvider
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.run import AgentRunResult
from haiku.rag.config import Config
from haiku.rag.research.dependencies import ResearchDependencies
@ -15,18 +14,18 @@ from haiku.rag.research.dependencies import ResearchDependencies
T = TypeVar("T")
class BaseResearchAgent(ABC):
class BaseResearchAgent(ABC, Generic[T]):
"""Base class for all research agents."""
def __init__(
self,
provider: str,
model: str,
output_type: type[T] | None = None,
output_type: type[T],
):
self.provider = provider
self.model = model
self.output_type = output_type or str
self.output_type = output_type
model_obj = self._get_model(provider, model)
@ -69,12 +68,14 @@ class BaseResearchAgent(ABC):
"""Register agent-specific tools."""
pass
async def run(self, prompt: str, deps: ResearchDependencies, **kwargs) -> Any:
async def run(
self, prompt: str, deps: ResearchDependencies, **kwargs
) -> AgentRunResult[T]:
"""Execute the agent."""
return await self._agent.run(prompt, deps=deps, **kwargs)
@property
def agent(self) -> Agent[ResearchDependencies, Any]:
def agent(self) -> Agent[ResearchDependencies, T]:
"""Access the underlying Pydantic AI agent."""
return self._agent

79
src/haiku/rag/research/clarification_agent.py
View file

@ -1,79 +0,0 @@
"""Clarification agent for gap detection and follow-up question generation."""
from pydantic import BaseModel, Field
from haiku.rag.research.base import BaseResearchAgent
class ClarificationResult(BaseModel):
"""Result of clarification analysis."""
information_gaps: list[str] = Field(
description="Specific missing information identified"
)
follow_up_questions: list[str] = Field(
description="Questions to ask to fill the gaps"
)
suggested_searches: list[str] = Field(
description="Recommended search queries for deeper investigation"
)
completeness_assessment: str = Field(
description="Overall assessment of research completeness"
)
priority_areas: list[str] = Field(
description="Most important areas to investigate next"
)
is_sufficient: bool = Field(
description="Whether the research has gathered sufficient information to answer the question"
)
confidence_score: float = Field(
ge=0.0,
le=1.0,
description="Confidence level (0-1) that the research is complete",
)
reasoning: str = Field(
description="Detailed reasoning for the completeness assessment"
)
class ClarificationAgent(BaseResearchAgent):
"""Agent specialized in identifying gaps and generating follow-up questions."""
def __init__(self, provider: str, model: str):
super().__init__(provider, model, output_type=ClarificationResult)
def get_system_prompt(self) -> str:
return """You are a clarification specialist agent focused on research completeness and quality.
Your role is to:
1. Critically evaluate what information has been gathered
2. Identify what crucial information is still missing
3. Detect contradictions or inconsistencies that need resolution
4. Generate targeted follow-up questions to fill knowledge gaps
5. Suggest specific search queries for deeper investigation
6. Assess the overall completeness of the research
Be thorough and critical in your evaluation. Consider:
- What questions remain unanswered?
- What assumptions need verification?
- What contradictions need resolution?
- What perspectives are missing?
- What details would strengthen the understanding?
IMPORTANT: When setting 'is_sufficient':
- True means: The research has enough information to provide a meaningful, accurate answer
- False means: Critical information is missing that prevents a complete answer
- Consider the nature of the question - simple questions need less, complex ones need more
- Be honest about uncertainty - if you're not confident, set is_sufficient to False
Your 'confidence_score' should reflect:
- 0.9-1.0: Very confident, all major aspects covered
- 0.7-0.9: Good coverage, minor gaps acceptable
- 0.5-0.7: Moderate coverage, some important gaps
- Below 0.5: Significant gaps, much more research needed"""
def register_tools(self) -> None:
"""Register clarification-specific tools."""
# The agent will use its LLM capabilities directly for gap analysis
# The structured output will guide the clarification process
pass

5
src/haiku/rag/research/dependencies.py
View file

@ -1,5 +1,3 @@
"""Shared dependencies for multi-agent research workflow."""
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field
@ -29,9 +27,6 @@ class ResearchContext(BaseModel):
gaps: list[str] = Field(
default_factory=list, description="Identified information gaps"
)
follow_up_questions: list[str] = Field(
default_factory=list, description="Generated follow-up questions"
)
def add_search_result(self, query: str, results: list["SearchResult"]) -> None:
"""Add search results to context."""

40
src/haiku/rag/research/evaluation_agent.py Normal file
View file

@ -0,0 +1,40 @@
from pydantic import BaseModel, Field
from haiku.rag.research.base import BaseResearchAgent
from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
class EvaluationResult(BaseModel):
"""Result of analysis and evaluation."""
key_insights: list[str] = Field(
description="Main insights extracted from the research so far"
)
new_questions: list[str] = Field(
description="New sub-questions to add to the research (max 3)", max_length=3
)
confidence_score: float = Field(
description="Confidence level in the completeness of research (0-1)",
ge=0.0,
le=1.0,
)
is_sufficient: bool = Field(
description="Whether the research is sufficient to answer the original question"
)
reasoning: str = Field(
description="Explanation of why the research is or isn't complete"
)
class AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]):
"""Agent that analyzes findings and evaluates research completeness."""
def __init__(self, provider: str, model: str) -> None:
super().__init__(provider, model, output_type=EvaluationResult)
def get_system_prompt(self) -> str:
return EVALUATION_AGENT_PROMPT
def register_tools(self) -> None:
"""No additional tools needed - uses LLM capabilities directly."""
pass

335
src/haiku/rag/research/orchestrator.py
View file

@ -1,17 +1,18 @@
from typing import Any
from pydantic import BaseModel, Field
from pydantic_ai import RunContext
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.analysis_agent import AnalysisAgent, AnalysisResult
from haiku.rag.research.base import BaseResearchAgent
from haiku.rag.research.clarification_agent import (
ClarificationAgent,
ClarificationResult,
)
from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.research.evaluation_agent import (
AnalysisEvaluationAgent,
EvaluationResult,
)
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
@ -21,11 +22,11 @@ class ResearchPlan(BaseModel):
main_question: str = Field(description="The main research question")
sub_questions: list[str] = Field(
description="Decomposed sub-questions to investigate"
description="Decomposed sub-questions to investigate (max 3)", max_length=3
)
class ResearchOrchestrator(BaseResearchAgent):
class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
"""Orchestrator agent that coordinates the research workflow."""
def __init__(
@ -37,108 +38,40 @@ class ResearchOrchestrator(BaseResearchAgent):
super().__init__(provider, model, output_type=ResearchPlan)
self.search_agent = SearchSpecialistAgent(provider, model)
self.analysis_agent = AnalysisAgent(provider, model)
self.clarification_agent = ClarificationAgent(provider, model)
self.synthesis_agent = SynthesisAgent(provider, model)
self.search_agent: SearchSpecialistAgent = SearchSpecialistAgent(
provider, model
)
self.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent(
provider, model
)
self.synthesis_agent: SynthesisAgent = SynthesisAgent(provider, model)
def get_system_prompt(self) -> str:
return """You are a research orchestrator responsible for coordinating a comprehensive research workflow.
Your role is to:
1. Understand and decompose the research question
2. Plan a systematic research approach
3. Coordinate specialized agents to gather and analyze information
4. Ensure comprehensive coverage of the topic
5. Iterate based on findings and gaps
Create a research plan that:
- Breaks down complex questions into manageable parts
- Identifies multiple search strategies
- Defines clear success criteria
- Ensures thorough investigation
/no_think"""
return ORCHESTRATOR_PROMPT
def register_tools(self) -> None:
"""Register orchestration tools."""
# Tools are no longer needed - orchestrator directly calls agents
pass
@self.agent.tool
async def delegate_search(
ctx: RunContext[ResearchDependencies], queries: list[str], limit: int = 5
) -> list[Any]:
"""Delegate search to the search specialist agent for multiple queries."""
all_results = []
# Search for each query
# The search agent will automatically store results in context
for query in queries:
result = await self.search_agent.run(
f"Search for: {query} with limit {limit}",
deps=ctx.deps,
usage=ctx.usage,
)
all_results.append(result)
return all_results
@self.agent.tool
async def delegate_analysis(
ctx: RunContext[ResearchDependencies],
) -> AnalysisResult:
"""Delegate analysis to the analysis agent."""
# Get search results from context
all_documents = []
for search in ctx.deps.context.search_results:
all_documents.extend(search.get("results", []))
# Pass documents for analysis
result = await self.analysis_agent.run(
f"Analyze these {len(all_documents)} documents from our search",
deps=ctx.deps,
usage=ctx.usage,
)
# Store analysis insights in context
if hasattr(result, "output") and isinstance(result.output, AnalysisResult):
for insight in result.output.key_insights:
ctx.deps.context.add_insight(insight)
return result.output if hasattr(result, "output") else result
@self.agent.tool
async def delegate_clarification(
ctx: RunContext[ResearchDependencies],
) -> ClarificationResult:
"""Delegate gap analysis to the clarification agent."""
result = await self.clarification_agent.run(
f"Evaluate the completeness of research on: {ctx.deps.context.original_question}",
deps=ctx.deps,
usage=ctx.usage,
)
# Store identified gaps in context
if hasattr(result, "output") and isinstance(
result.output, ClarificationResult
):
for gap in result.output.information_gaps:
ctx.deps.context.add_gap(gap)
ctx.deps.context.follow_up_questions.extend(
result.output.follow_up_questions
)
return result.output if hasattr(result, "output") else result
@self.agent.tool
async def generate_report(
ctx: RunContext[ResearchDependencies],
) -> ResearchReport:
"""Generate final research report using synthesis agent."""
result = await self.synthesis_agent.run(
f"Create a comprehensive research report for: {ctx.deps.context.original_question}",
deps=ctx.deps,
usage=ctx.usage,
)
return result.output if hasattr(result, "output") else result
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["question"],
"answer": qa["answer"][:500] + "..."
if len(qa["answer"]) > 500
else qa["answer"],
}
for qa in context.qa_responses
],
"insights": context.insights,
"gaps": context.gaps,
}
return format_as_xml(context_data, root_tag="research_context")
async def conduct_research(
self,
@ -174,11 +107,10 @@ class ResearchOrchestrator(BaseResearchAgent):
if console:
console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
plan_result = await self.run(
plan_result: AgentRunResult[ResearchPlan] = await self.run(
f"Create a research plan for: {question}", deps=deps
)
assert plan_result.output and isinstance(plan_result.output, ResearchPlan)
context.sub_questions = plan_result.output.sub_questions
if console:
@ -198,99 +130,101 @@ class ResearchOrchestrator(BaseResearchAgent):
f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]"
)
# Determine what to search for in this iteration
if context.follow_up_questions:
# Use follow-up questions from previous clarification
search_target = context.follow_up_questions[:3] # Take top 3 follow-ups
search_prompt = ", ".join(search_target)
elif iteration < len(context.sub_questions):
# Use pre-planned sub-questions
search_prompt = context.sub_questions[iteration]
else:
# Fall back to original question
search_prompt = question
# Search phase - directly call the search agent
if console:
console.print(
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {search_prompt}"
)
await self.search_agent.run(search_prompt, deps=deps)
if console:
# Show documents found
if context.search_results:
latest_results = context.search_results[-1]
console.print(
f" Found [green]{len(latest_results.get('results', []))} documents[/green]"
)
# Show the answer generated
if context.qa_responses:
latest_qa = context.qa_responses[-1]
answer_preview = (
latest_qa["answer"][:200] + "..."
if len(latest_qa["answer"]) > 200
else latest_qa["answer"]
)
console.print(f" [bold]Answer:[/bold] {answer_preview}")
# Analysis phase (only if we have results)
if context.search_results:
# Check if we have questions to search
if not context.sub_questions:
# No more questions to explore
if console:
console.print(
"\n[bold cyan]📊 Analyzing gathered information...[/bold cyan]"
"[yellow]No more questions to explore. Concluding research.[/yellow]"
)
break
analysis_result = await self.analysis_agent.run(
"Analyze the gathered information", deps=deps
)
# Use current sub-questions for this iteration
questions_to_search = context.sub_questions
if console and hasattr(analysis_result, "output"):
output = analysis_result.output
if hasattr(output, "key_insights") and output.key_insights:
console.print(" [bold]Key insights:[/bold]")
for insight in output.key_insights[:3]:
console.print(f" • {insight}")
# Clarification phase - evaluate completeness
# Search phase - answer all questions in this iteration
if console:
console.print(
"\n[bold cyan]🔎 Evaluating research completeness...[/bold cyan]"
f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]"
)
for i, q in enumerate(questions_to_search, 1):
console.print(f" {i}. {q}")
# Run searches for all questions and remove answered ones
answered_questions = []
for search_question in questions_to_search:
await self.search_agent.run(search_question, deps=deps)
# Mark this question as answered
answered_questions.append(search_question)
if 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"]
)
console.print(
f"\n [green]✓[/green] {search_question[:50]}..."
if len(search_question) > 50
else f"\n [green]✓[/green] {search_question}"
)
console.print(f" {answer_preview}")
# Remove answered questions from the list
for question in answered_questions:
if question in context.sub_questions:
context.sub_questions.remove(question)
# Analysis and Evaluation phase
if console:
console.print(
"\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]"
)
clarification_result = await self.clarification_agent.run(
f"Evaluate the completeness of research for: {question}. "
f"Consider all information gathered so far and determine if we have sufficient "
f"information to provide a comprehensive answer.",
# Format context for the evaluation agent
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,
)
if console and hasattr(clarification_result, "output"):
output = clarification_result.output
if hasattr(output, "confidence_score"):
console.print(
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]"
)
if hasattr(output, "is_sufficient"):
status = (
"[green]Yes[/green]"
if output.is_sufficient
else "[red]No[/red]"
)
console.print(f" Sufficient: {status}")
if console and evaluation_result.output:
output = evaluation_result.output
if output.key_insights:
console.print(" [bold]Key insights:[/bold]")
for insight in output.key_insights[:3]:
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}")
# Check if research is sufficient based on semantic evaluation
if self._should_stop_research(clarification_result, confidence_threshold):
# Log the reasoning for stopping
if (
console
and hasattr(clarification_result, "output")
and isinstance(clarification_result.output, ClarificationResult)
):
# 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):
if console:
console.print(
f"\n[bold green]✅ Stopping research:[/bold green] {clarification_result.output.reasoning}"
f"\n[bold green]✅ Stopping research:[/bold green] {evaluation_result.output.reasoning}"
)
break
@ -300,30 +234,31 @@ class ResearchOrchestrator(BaseResearchAgent):
"\n[bold cyan]📝 Generating final research report...[/bold cyan]"
)
report_result = await self.synthesis_agent.run(
"Generate the final research report", deps=deps
# 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
)
if console:
console.print("[bold green]✅ Research complete![/bold green]")
return (
report_result.output if hasattr(report_result, "output") else report_result
)
return report_result.output
def _should_stop_research(
self, clarification_result: Any, confidence_threshold: float
self,
evaluation_result: AgentRunResult[EvaluationResult],
confidence_threshold: float,
) -> bool:
"""Determine if research should stop based on semantic completeness evaluation."""
"""Determine if research should stop based on evaluation."""
if not hasattr(clarification_result, "output") or not isinstance(
clarification_result.output, ClarificationResult
):
# If we can't evaluate, continue researching
return False
result = evaluation_result.output
result = clarification_result.output
# Use the LLM's semantic evaluation
# Stop if the agent indicates sufficient information AND confidence exceeds threshold
return result.is_sufficient and result.confidence_score >= confidence_threshold

87
src/haiku/rag/research/prompts.py Normal file
View file

@ -0,0 +1,87 @@
ORCHESTRATOR_PROMPT = """You are a research orchestrator responsible for coordinating a comprehensive research workflow.
Your role is to:
1. Understand and decompose the research question
2. Plan a systematic research approach
3. Coordinate specialized agents to gather and analyze information
4. Ensure comprehensive coverage of the topic
5. Iterate based on findings and gaps
Create a research plan that:
- Breaks down the question into at most 3 focused sub-questions
- Each sub-question should target a specific aspect of the research
- Prioritize the most important aspects to investigate
- Ensure comprehensive coverage within the 3-question limit"""
SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist.
Your role is to:
1. Search the knowledge base for relevant information
2. Analyze the retrieved documents
3. Provide a comprehensive answer to the question
4. Base your answer strictly on the information found
Use the search_and_answer tool to retrieve relevant documents and formulate your response.
Be thorough and specific in your answers, citing relevant information from the sources."""
EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for research workflows.
You have access to:
- The original research question
- Question-answer pairs from search operations
- Raw search results and source documents
- Previously identified insights
Your dual role is to:
ANALYSIS:
1. Extract key insights from all gathered information
2. Identify patterns and connections across sources
3. Synthesize findings into coherent understanding
4. Focus on the most important discoveries
EVALUATION:
1. Assess if we have sufficient information to answer the original question
2. Calculate a confidence score (0-1) based on:
- Coverage of the main question's aspects
- Quality and consistency of sources
- Depth of information gathered
3. Identify specific gaps that still need investigation
4. Generate up to 3 new sub-questions that haven't been answered yet
Be critical and thorough in your evaluation. Only mark research as sufficient when:
- All major aspects of the question are addressed
- Sources provide consistent, reliable information
- The depth of coverage meets the question's requirements
- No critical gaps remain
Generate new sub-questions that:
- Target specific unexplored aspects not covered by existing questions
- Seek clarification on ambiguities
- Explore important edge cases or exceptions
- Are focused and actionable (max 3)
- Do NOT repeat or rephrase questions that have already been answered (see qa_responses)
- Should be genuinely new areas to explore"""
SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist agent focused on creating comprehensive research reports.
Your role is to:
1. Synthesize all gathered information into a coherent narrative
2. Present findings in a clear, structured format
3. Draw evidence-based conclusions
4. Acknowledge limitations and uncertainties
5. Provide actionable recommendations
6. Maintain academic rigor and objectivity
Your report should be:
- Comprehensive yet concise
- Well-structured and easy to follow
- Based solely on evidence from the research
- Transparent about limitations
- Professional and objective in tone
Focus on creating a report that provides clear value to the reader by:
- Answering the original research question thoroughly
- Highlighting the most important findings
- Explaining the implications of the research
- Suggesting concrete next steps"""

27
src/haiku/rag/research/search_agent.py
View file

@ -1,27 +1,27 @@
"""Search specialist agent that answers questions using RAG."""
from typing import Any
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 SEARCH_AGENT_PROMPT
class SearchSpecialistAgent(BaseResearchAgent):
class SearchSpecialistAgent(BaseResearchAgent[str]):
"""Agent specialized in answering questions using RAG search."""
def __init__(self, provider: str, model: str):
def __init__(self, provider: str, model: str) -> None:
# Output is a string answer, like the QA agent
super().__init__(provider, model, output_type=str)
async def run(self, prompt: str, deps: ResearchDependencies, **kwargs) -> Any:
async def run(
self, prompt: str, deps: ResearchDependencies, **kwargs
) -> AgentRunResult[str]:
"""Execute the agent and store QA response in context."""
# Run the base agent
result = await super().run(prompt, deps, **kwargs)
# Store the QA response if we got an answer
if hasattr(result, "output") and result.output:
if result.output:
# Get the sources from the last search (which the tool just stored)
if deps.context.search_results:
last_search = deps.context.search_results[-1]
@ -31,16 +31,7 @@ class SearchSpecialistAgent(BaseResearchAgent):
return result
def get_system_prompt(self) -> str:
return """You are a search and question-answering specialist.
Your role is to:
1. Search the knowledge base for relevant information
2. Analyze the retrieved documents
3. Provide a comprehensive answer to the question
4. Base your answer strictly on the information found
Use the search_and_answer tool to retrieve relevant documents and formulate your response.
Be thorough and specific in your answers, citing relevant information from the sources."""
return SEARCH_AGENT_PROMPT
def register_tools(self) -> None:
"""Register search-specific tools."""

30
src/haiku/rag/research/synthesis_agent.py
View file

@ -1,8 +1,7 @@
"""Synthesis agent for final research report generation."""
from pydantic import BaseModel, Field
from haiku.rag.research.base import BaseResearchAgent
from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT
class ResearchReport(BaseModel):
@ -24,35 +23,14 @@ class ResearchReport(BaseModel):
)
class SynthesisAgent(BaseResearchAgent):
class SynthesisAgent(BaseResearchAgent[ResearchReport]):
"""Agent specialized in synthesizing research into comprehensive reports."""
def __init__(self, provider: str, model: str):
def __init__(self, provider: str, model: str) -> None:
super().__init__(provider, model, output_type=ResearchReport)
def get_system_prompt(self) -> str:
return """You are a synthesis specialist agent focused on creating comprehensive research reports.
Your role is to:
1. Synthesize all gathered information into a coherent narrative
2. Present findings in a clear, structured format
3. Draw evidence-based conclusions
4. Acknowledge limitations and uncertainties
5. Provide actionable recommendations
6. Maintain academic rigor and objectivity
Your report should be:
- Comprehensive yet concise
- Well-structured and easy to follow
- Based solely on evidence from the research
- Transparent about limitations
- Professional and objective in tone
Focus on creating a report that provides clear value to the reader by:
- Answering the original research question thoroughly
- Highlighting the most important findings
- Explaining the implications of the research
- Suggesting concrete next steps"""
return SYNTHESIS_AGENT_PROMPT
def register_tools(self) -> None:
"""Register synthesis-specific tools."""

237
tests/research/test_orchestrator.py
View file

@ -3,15 +3,10 @@
from unittest.mock import AsyncMock, MagicMock, create_autospec
import pytest
from pydantic_ai import RunContext
from pydantic_ai.models.test import TestModel
from pydantic_ai.usage import RunUsage
from haiku.rag.client import HaikuRAG
from haiku.rag.research.analysis_agent import AnalysisResult
from haiku.rag.research.base import SearchResult
from haiku.rag.research.clarification_agent import ClarificationResult
from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.research.evaluation_agent import EvaluationResult
from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan
from haiku.rag.store.models.chunk import Chunk
@ -62,8 +57,10 @@ class TestResearchOrchestrator:
# All agents should use the same provider/model
assert orchestrator.search_agent.provider == orchestrator.provider
assert orchestrator.search_agent.model == orchestrator.model
assert orchestrator.analysis_agent.provider == orchestrator.provider
assert orchestrator.analysis_agent.model == orchestrator.model
assert orchestrator.evaluation_agent.provider == orchestrator.provider
assert orchestrator.evaluation_agent.model == orchestrator.model
assert orchestrator.synthesis_agent.provider == orchestrator.provider
assert orchestrator.synthesis_agent.model == orchestrator.model
def test_orchestrator_initialization(self):
"""Test that orchestrator initializes all agents correctly."""
@ -71,63 +68,57 @@ class TestResearchOrchestrator:
# Check all agents are initialized
assert orchestrator.search_agent is not None
assert orchestrator.analysis_agent is not None
assert orchestrator.clarification_agent is not None
assert orchestrator.evaluation_agent is not None
assert orchestrator.synthesis_agent is not None
# Check they all use the same provider and model
assert orchestrator.search_agent.provider == "openai"
assert orchestrator.search_agent.model == "gpt-4"
assert orchestrator.analysis_agent.provider == "openai"
assert orchestrator.clarification_agent.provider == "openai"
assert orchestrator.evaluation_agent.provider == "openai"
assert orchestrator.evaluation_agent.model == "gpt-4"
assert orchestrator.synthesis_agent.provider == "openai"
assert orchestrator.synthesis_agent.model == "gpt-4"
def test_orchestrator_has_correct_output_type(self):
"""Test that orchestrator's output type is ResearchPlan."""
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
assert orchestrator.output_type == ResearchPlan
def test_orchestrator_registers_delegation_tools(self):
"""Test that orchestrator registers all delegation tools."""
def test_orchestrator_has_no_tools(self):
"""Test that orchestrator no longer registers tools (direct agent calls now)."""
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
# Get the tools from the agent
tools = orchestrator.agent._function_toolset.tools
tool_names = list(tools.keys())
# Check all delegation tools are registered
assert "delegate_search" in tool_names
assert "delegate_analysis" in tool_names
assert "delegate_clarification" in tool_names
assert "generate_report" in tool_names
# Should have no tools since we call agents directly now
assert len(tool_names) == 0
def test_should_stop_research_logic(self):
"""Test the stopping logic based on ClarificationResult."""
"""Test the stopping logic based on EvaluationResult."""
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
# Create mock clarification results
# Create mock evaluation results
sufficient_result = MagicMock()
sufficient_result.output = ClarificationResult(
information_gaps=[],
follow_up_questions=[],
suggested_searches=[],
completeness_assessment="Research is comprehensive",
priority_areas=[],
is_sufficient=True,
sufficient_result.output = EvaluationResult(
key_insights=["Climate is changing", "Human activity is the cause"],
new_questions=[],
confidence_score=0.9,
reasoning="All aspects covered",
is_sufficient=True,
reasoning="All aspects covered comprehensively",
)
insufficient_result = MagicMock()
insufficient_result.output = ClarificationResult(
information_gaps=["Missing data on impacts"],
follow_up_questions=["What about economic impacts?"],
suggested_searches=["economic impact climate change"],
completeness_assessment="More research needed",
priority_areas=["Economic analysis"],
is_sufficient=False,
insufficient_result.output = EvaluationResult(
key_insights=["Some data found"],
new_questions=[
"What about economic impacts?",
"Regional variations?",
],
confidence_score=0.4,
reasoning="Major gaps remain",
is_sufficient=False,
reasoning="Major gaps remain in understanding",
)
# Test with sufficient research (threshold 0.8)
@ -145,138 +136,66 @@ class TestResearchOrchestrator:
assert not orchestrator._should_stop_research(insufficient_result, 0.8)
@pytest.mark.asyncio
async def test_delegate_search_tool(self, research_deps):
"""Test the delegate_search tool function."""
async def test_conduct_research_workflow(self, mock_client):
"""Test the basic research workflow."""
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
# Get the delegate_search tool
tools = orchestrator.agent._function_toolset.tools
delegate_search = tools["delegate_search"].function
# Mock the search agent's run method
orchestrator.search_agent.run = AsyncMock(
return_value=MagicMock(output=["results"])
# Mock the agent runs
# Mock initial plan
plan_mock = MagicMock()
plan_mock.output = ResearchPlan(
main_question="What is climate change?",
sub_questions=[
"What causes climate change?",
"What are the effects?",
"What can be done?",
],
)
orchestrator.run = AsyncMock(return_value=plan_mock)
# Create context and call the tool
ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage())
await delegate_search(ctx, queries=["climate change"])
# Mock search agent
search_mock = MagicMock()
search_mock.output = "Climate change is caused by greenhouse gases."
orchestrator.search_agent.run = AsyncMock(return_value=search_mock)
# Verify the search agent was called
orchestrator.search_agent.run.assert_called_once()
assert "climate change" in orchestrator.search_agent.run.call_args[0][0]
@pytest.mark.asyncio
async def test_delegate_analysis_tool(self, research_deps):
"""Test the delegate_analysis tool function."""
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
# Add some search results to context
research_deps.context.search_results = [
{
"query": "test",
"results": [
SearchResult(
content="Climate data",
score=0.9,
document_uri="doc1.md",
metadata={},
)
],
}
]
# Get the delegate_analysis tool
tools = orchestrator.agent._function_toolset.tools
delegate_analysis = tools["delegate_analysis"].function
# Mock the analysis agent's run method
mock_result = MagicMock()
mock_result.output = AnalysisResult(
key_insights=["Climate is changing"],
themes={"warming": ["temperature rise"]},
summary="Analysis complete",
evidence_quality="strong",
recommendations=["More research needed"],
# Mock evaluation agent - make it stop after first iteration
eval_mock = MagicMock()
eval_mock.output = EvaluationResult(
key_insights=["Climate change is real"],
new_questions=[],
confidence_score=0.9,
is_sufficient=True,
reasoning="Sufficient information gathered",
)
orchestrator.analysis_agent.run = AsyncMock(return_value=mock_result)
orchestrator.evaluation_agent.run = AsyncMock(return_value=eval_mock)
# Create context and call the tool
ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage())
await delegate_analysis(ctx)
# Verify the analysis agent was called
orchestrator.analysis_agent.run.assert_called_once()
# Verify insights were added to context
assert "Climate is changing" in research_deps.context.insights
@pytest.mark.asyncio
async def test_delegate_clarification_tool(self, research_deps):
"""Test the delegate_clarification tool function."""
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
# Get the delegate_clarification tool
tools = orchestrator.agent._function_toolset.tools
delegate_clarification = tools["delegate_clarification"].function
# Mock the clarification agent's run method
mock_result = MagicMock()
mock_result.output = ClarificationResult(
information_gaps=["Missing economic data"],
follow_up_questions=["What about costs?"],
suggested_searches=["climate change costs"],
completeness_assessment="Needs more data",
priority_areas=["Economics"],
is_sufficient=False,
confidence_score=0.6,
reasoning="Missing key information",
)
orchestrator.clarification_agent.run = AsyncMock(return_value=mock_result)
# Create context and call the tool
ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage())
await delegate_clarification(ctx)
# Verify the clarification agent was called
orchestrator.clarification_agent.run.assert_called_once()
# Verify gaps and questions were added to context
assert "Missing economic data" in research_deps.context.gaps
assert "What about costs?" in research_deps.context.follow_up_questions
@pytest.mark.asyncio
async def test_generate_report_tool(self, research_deps):
"""Test the generate_report tool function."""
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
# Get the generate_report tool
tools = orchestrator.agent._function_toolset.tools
generate_report = tools["generate_report"].function
# Mock the synthesis agent's run method
# Mock synthesis agent
from haiku.rag.research.synthesis_agent import ResearchReport
mock_result = MagicMock()
mock_result.output = ResearchReport(
title="Climate Change Research",
executive_summary="Summary of findings",
main_findings=["Finding 1", "Finding 2"],
themes={"warming": "Global temperature rise"},
conclusions=["Conclusion 1"],
limitations=["Limited data"],
recommendations=["More research"],
sources_summary="Various sources",
synthesis_mock = MagicMock()
synthesis_mock.output = ResearchReport(
title="Climate Change Report",
executive_summary="Summary",
main_findings=["Finding 1"],
themes={},
conclusions=[],
limitations=[],
recommendations=[],
sources_summary="Sources",
)
orchestrator.synthesis_agent.run = AsyncMock(return_value=mock_result)
orchestrator.synthesis_agent.run = AsyncMock(return_value=synthesis_mock)
# Create context and call the tool
ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage())
result = await generate_report(ctx)
# Mock client search and expand
mock_client.search.return_value = []
mock_client.expand_context.return_value = []
# Verify the synthesis agent was called
orchestrator.synthesis_agent.run.assert_called_once()
# Run the research
report = await orchestrator.conduct_research(
"What is climate change?", mock_client, max_iterations=3
)
# Verify we got a ResearchReport
assert isinstance(result, ResearchReport)
assert result.title == "Climate Change Research"
# Verify we got a report
assert report.title == "Climate Change Report"
# Verify search was called for all 3 sub-questions
assert orchestrator.search_agent.run.call_count == 3