from typing import Any from pydantic import BaseModel, Field from pydantic_ai import RunContext 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.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" ) search_strategies: list[str] = Field( description="Different search approaches to use" ) success_criteria: list[str] = Field(description="Criteria for successful research") class ResearchOrchestrator(BaseResearchAgent): """Orchestrator agent that coordinates the research workflow.""" def __init__( self, provider: str = Config.RERANK_PROVIDER, model: str = Config.RERANK_MODEL ): 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) 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""" def register_tools(self) -> None: """Register orchestration tools.""" @self.agent.tool async def delegate_search( ctx: RunContext[ResearchDependencies], queries: list[str], limit: int = 5 ) -> Any: """Delegate search to the search specialist agent.""" # Pass the context to maintain usage tracking result = await self.search_agent.run( f"Search for: {', '.join(queries)}", deps=ctx.deps, usage=ctx.usage ) return result @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 async def conduct_research( self, question: str, client: Any, max_iterations: int = 3, confidence_threshold: float = 0.8, ) -> 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) Returns: ResearchReport with comprehensive findings """ # Initialize context context = ResearchContext(original_question=question) deps = ResearchDependencies(client=client, context=context) # Create initial research plan plan_result = await self.run( f"Create a research plan for: {question}", deps=deps ) if hasattr(plan_result, "output") and isinstance( plan_result.output, ResearchPlan ): context.sub_questions = plan_result.output.sub_questions # Execute research iterations for iteration in range(max_iterations): # 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 = f"Search for: {', '.join(search_target)}" elif iteration < len(context.sub_questions): # Use pre-planned sub-questions search_prompt = f"Search for: {context.sub_questions[iteration]}" else: # Fall back to original question with variation search_prompt = f"Additional search for: {question}" # Search phase await self.run(search_prompt, deps=deps) # Analysis phase (only if we have results) if context.search_results: await self.run("Analyze the gathered information", deps=deps) # Clarification phase - evaluate completeness clarification_result = await self.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.", deps=deps, ) # Check if research is sufficient based on semantic evaluation if self._should_stop_research(clarification_result, confidence_threshold): # Log the reasoning for stopping if hasattr(clarification_result, "output") and isinstance( clarification_result.output, ClarificationResult ): print(f"Stopping research: {clarification_result.output.reasoning}") break # Generate final report report_result = await self.run("Generate the final research report", deps=deps) return ( report_result.output if hasattr(report_result, "output") else report_result ) def _should_stop_research( self, clarification_result: Any, confidence_threshold: float ) -> bool: """Determine if research should stop based on semantic completeness evaluation.""" if not hasattr(clarification_result, "output") or not isinstance( clarification_result.output, ClarificationResult ): # If we can't evaluate, continue researching return False 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