265 lines
10 KiB
Python
265 lines
10 KiB
Python
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.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
|
||
|
||
|
||
class ResearchPlan(BaseModel):
|
||
"""Research execution plan."""
|
||
|
||
main_question: str = Field(description="The main research question")
|
||
sub_questions: list[str] = Field(
|
||
description="Decomposed sub-questions to investigate (max 3)", max_length=3
|
||
)
|
||
|
||
|
||
class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
|
||
"""Orchestrator agent that coordinates the research workflow."""
|
||
|
||
def __init__(
|
||
self, provider: str | None = Config.RESEARCH_PROVIDER, model: str | None = None
|
||
):
|
||
# Use provided values or fall back to config defaults
|
||
provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER
|
||
model = model or Config.RESEARCH_MODEL or Config.QA_MODEL
|
||
|
||
super().__init__(provider, model, output_type=ResearchPlan)
|
||
|
||
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 ORCHESTRATOR_PROMPT
|
||
|
||
def register_tools(self) -> None:
|
||
"""Register orchestration tools."""
|
||
# Tools are no longer needed - orchestrator directly calls agents
|
||
pass
|
||
|
||
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}
|
||
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,
|
||
question: str,
|
||
client: Any,
|
||
max_iterations: int = 3,
|
||
confidence_threshold: float = 0.8,
|
||
verbose: bool = False,
|
||
console: Console | None = None,
|
||
) -> ResearchReport:
|
||
"""Conduct comprehensive research on a question.
|
||
|
||
Args:
|
||
question: The research question to investigate
|
||
client: HaikuRAG client for document operations
|
||
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
|
||
"""
|
||
|
||
# Initialize context
|
||
context = ResearchContext(original_question=question)
|
||
deps = ResearchDependencies(client=client, context=context)
|
||
|
||
# Use provided console or create a new one
|
||
console = console or Console() if verbose else None
|
||
|
||
# Create initial research plan
|
||
if console:
|
||
console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
|
||
|
||
plan_result: AgentRunResult[ResearchPlan] = await self.run(
|
||
f"Create a research plan for: {question}", deps=deps
|
||
)
|
||
|
||
context.sub_questions = plan_result.output.sub_questions
|
||
|
||
if console:
|
||
console.print("\n[bold green]✅ Research Plan Created:[/bold green]")
|
||
console.print(
|
||
f" [bold]Main Question:[/bold] {plan_result.output.main_question}"
|
||
)
|
||
console.print(" [bold]Sub-questions:[/bold]")
|
||
for i, sq in enumerate(plan_result.output.sub_questions, 1):
|
||
console.print(f" {i}. {sq}")
|
||
console.print()
|
||
|
||
# Execute research iterations
|
||
for iteration in range(max_iterations):
|
||
if console:
|
||
console.rule(
|
||
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:
|
||
console.print(
|
||
"[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
|
||
if console:
|
||
console.print(
|
||
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:
|
||
try:
|
||
await self.search_agent.run(search_question, deps=deps)
|
||
except Exception as e: # pragma: no cover - defensive
|
||
if console:
|
||
console.print(
|
||
f"\n [red]×[/red] Omitting failed question: {search_question} ({e})"
|
||
)
|
||
finally:
|
||
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]"
|
||
)
|
||
|
||
# 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 evaluation_result.output:
|
||
output = evaluation_result.output
|
||
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}")
|
||
|
||
# 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] {evaluation_result.output.reasoning}"
|
||
)
|
||
break
|
||
|
||
# Generate final report
|
||
if console:
|
||
console.print(
|
||
"\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
|
||
)
|
||
|
||
if console:
|
||
console.print("[bold green]✅ Research complete![/bold green]")
|
||
|
||
return report_result.output
|
||
|
||
def _should_stop_research(
|
||
self,
|
||
evaluation_result: AgentRunResult[EvaluationResult],
|
||
confidence_threshold: float,
|
||
) -> bool:
|
||
"""Determine if research should stop based on evaluation."""
|
||
|
||
result = evaluation_result.output
|
||
|
||
# Stop if the agent indicates sufficient information AND confidence exceeds threshold
|
||
return result.is_sufficient and result.confidence_score >= confidence_threshold
|