From 98ca5c4746cb26d3a35c5d583df629a49cae4d09 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 16 Sep 2025 14:25:28 +0300 Subject: [PATCH] Simplify orchestration & subagents --- src/haiku/rag/qa/agent.py | 6 +- src/haiku/rag/qa/prompts.py | 4 +- src/haiku/rag/research/__init__.py | 15 +- src/haiku/rag/research/analysis_agent.py | 53 --- src/haiku/rag/research/base.py | 17 +- src/haiku/rag/research/clarification_agent.py | 79 ----- src/haiku/rag/research/dependencies.py | 5 - src/haiku/rag/research/evaluation_agent.py | 40 +++ src/haiku/rag/research/orchestrator.py | 335 +++++++----------- src/haiku/rag/research/prompts.py | 87 +++++ src/haiku/rag/research/search_agent.py | 27 +- src/haiku/rag/research/synthesis_agent.py | 30 +- tests/research/test_orchestrator.py | 237 ++++--------- 13 files changed, 374 insertions(+), 561 deletions(-) delete mode 100644 src/haiku/rag/research/analysis_agent.py delete mode 100644 src/haiku/rag/research/clarification_agent.py create mode 100644 src/haiku/rag/research/evaluation_agent.py create mode 100644 src/haiku/rag/research/prompts.py diff --git a/src/haiku/rag/qa/agent.py b/src/haiku/rag/qa/agent.py index 594fad6e..b7edb6fa 100644 --- a/src/haiku/rag/qa/agent.py +++ b/src/haiku/rag/qa/agent.py @@ -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( diff --git a/src/haiku/rag/qa/prompts.py b/src/haiku/rag/qa/prompts.py index f865b09f..19896d2c 100644 --- a/src/haiku/rag/qa/prompts.py +++ b/src/haiku/rag/qa/prompts.py @@ -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. diff --git a/src/haiku/rag/research/__init__.py b/src/haiku/rag/research/__init__.py index 18e71cdd..d947df6e 100644 --- a/src/haiku/rag/research/__init__.py +++ b/src/haiku/rag/research/__init__.py @@ -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 diff --git a/src/haiku/rag/research/analysis_agent.py b/src/haiku/rag/research/analysis_agent.py deleted file mode 100644 index dc817bd9..00000000 --- a/src/haiku/rag/research/analysis_agent.py +++ /dev/null @@ -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 diff --git a/src/haiku/rag/research/base.py b/src/haiku/rag/research/base.py index 2cb5c1aa..8098fc7e 100644 --- a/src/haiku/rag/research/base.py +++ b/src/haiku/rag/research/base.py @@ -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 diff --git a/src/haiku/rag/research/clarification_agent.py b/src/haiku/rag/research/clarification_agent.py deleted file mode 100644 index 07ea3dfd..00000000 --- a/src/haiku/rag/research/clarification_agent.py +++ /dev/null @@ -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 diff --git a/src/haiku/rag/research/dependencies.py b/src/haiku/rag/research/dependencies.py index fa90eedb..bb64bf1b 100644 --- a/src/haiku/rag/research/dependencies.py +++ b/src/haiku/rag/research/dependencies.py @@ -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.""" diff --git a/src/haiku/rag/research/evaluation_agent.py b/src/haiku/rag/research/evaluation_agent.py new file mode 100644 index 00000000..5b924c81 --- /dev/null +++ b/src/haiku/rag/research/evaluation_agent.py @@ -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 diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index f73db6ec..278acc05 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -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 diff --git a/src/haiku/rag/research/prompts.py b/src/haiku/rag/research/prompts.py new file mode 100644 index 00000000..a4b36d91 --- /dev/null +++ b/src/haiku/rag/research/prompts.py @@ -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""" diff --git a/src/haiku/rag/research/search_agent.py b/src/haiku/rag/research/search_agent.py index d06b3830..82e937c4 100644 --- a/src/haiku/rag/research/search_agent.py +++ b/src/haiku/rag/research/search_agent.py @@ -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.""" diff --git a/src/haiku/rag/research/synthesis_agent.py b/src/haiku/rag/research/synthesis_agent.py index bb8928c8..45349789 100644 --- a/src/haiku/rag/research/synthesis_agent.py +++ b/src/haiku/rag/research/synthesis_agent.py @@ -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.""" diff --git a/tests/research/test_orchestrator.py b/tests/research/test_orchestrator.py index 6dce457e..635bbaea 100644 --- a/tests/research/test_orchestrator.py +++ b/tests/research/test_orchestrator.py @@ -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