diff --git a/src/haiku/rag/research/__init__.py b/src/haiku/rag/research/__init__.py index 29ac92c0..18e71cdd 100644 --- a/src/haiku/rag/research/__init__.py +++ b/src/haiku/rag/research/__init__.py @@ -1,6 +1,32 @@ """Multi-agent research workflow for advanced RAG queries.""" -from haiku.rag.research.base import BaseResearchAgent -from haiku.rag.research.dependencies import ResearchDependencies +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.orchestrator import ResearchOrchestrator, ResearchPlan +from haiku.rag.research.search_agent import SearchSpecialistAgent +from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent -__all__ = ["ResearchDependencies", "BaseResearchAgent"] +__all__ = [ + # Base classes + "BaseResearchAgent", + "ResearchDependencies", + "ResearchContext", + "SearchResult", + "ResearchOutput", + # Specialized agents + "SearchSpecialistAgent", + "AnalysisAgent", + "AnalysisResult", + "ClarificationAgent", + "ClarificationResult", + "SynthesisAgent", + "ResearchReport", + # Orchestrator + "ResearchOrchestrator", + "ResearchPlan", +] diff --git a/src/haiku/rag/research/clarification_agent.py b/src/haiku/rag/research/clarification_agent.py index e4aec7b3..07ea3dfd 100644 --- a/src/haiku/rag/research/clarification_agent.py +++ b/src/haiku/rag/research/clarification_agent.py @@ -23,6 +23,17 @@ class ClarificationResult(BaseModel): 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): @@ -49,7 +60,17 @@ class ClarificationAgent(BaseResearchAgent): - What perspectives are missing? - What details would strengthen the understanding? - Your goal is to ensure comprehensive, accurate, and complete research.""" + 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.""" diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py new file mode 100644 index 00000000..cdb748fa --- /dev/null +++ b/src/haiku/rag/research/orchestrator.py @@ -0,0 +1,225 @@ +"""Research orchestrator agent that coordinates specialized agents.""" + +from typing import Any + +from pydantic import BaseModel, Field +from pydantic_ai import RunContext + +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, model: str): + super().__init__(provider, model, output_type=ResearchPlan) + + # Initialize specialized agents + 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 diff --git a/tests/research/test_orchestrator.py b/tests/research/test_orchestrator.py new file mode 100644 index 00000000..1eccffa3 --- /dev/null +++ b/tests/research/test_orchestrator.py @@ -0,0 +1,268 @@ +"""Tests for the research orchestrator.""" + +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.orchestrator import ResearchOrchestrator, ResearchPlan +from haiku.rag.store.models.chunk import Chunk + + +@pytest.fixture +def mock_client(): + """Create a mock HaikuRAG client.""" + client = create_autospec(HaikuRAG, instance=True) + client.search = AsyncMock() + client.expand_context = AsyncMock() + return client + + +@pytest.fixture +def research_context(): + """Create a research context.""" + return ResearchContext(original_question="What is climate change?") + + +@pytest.fixture +def research_deps(mock_client, research_context): + """Create research dependencies.""" + return ResearchDependencies(client=mock_client, context=research_context) + + +def create_mock_chunk(chunk_id: str, content: str, score: float = 0.8): + """Helper to create mock chunk objects.""" + return Chunk( + id=chunk_id, + document_id=f"doc_{chunk_id}", + content=content, + document_uri=f"doc_{chunk_id}.md", + metadata={}, + ), score + + +class TestResearchOrchestrator: + """Test suite for ResearchOrchestrator.""" + + def test_orchestrator_initialization(self): + """Test that orchestrator initializes all agents correctly.""" + orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4") + + # 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.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.synthesis_agent.provider == "openai" + + 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.""" + 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 + + def test_should_stop_research_logic(self): + """Test the stopping logic based on ClarificationResult.""" + orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4") + + # Create mock clarification 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, + confidence_score=0.9, + reasoning="All aspects covered", + ) + + 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, + confidence_score=0.4, + reasoning="Major gaps remain", + ) + + # Test with sufficient research (threshold 0.8) + assert orchestrator._should_stop_research(sufficient_result, 0.8) + + # Test with insufficient research + assert not orchestrator._should_stop_research(insufficient_result, 0.8) + + # Test with high confidence but below threshold + sufficient_result.output.confidence_score = 0.75 + assert not orchestrator._should_stop_research(sufficient_result, 0.8) + + # Test with is_sufficient=False even with high confidence + insufficient_result.output.confidence_score = 0.95 + 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.""" + 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"]) + ) + + # Create context and call the tool + ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage()) + await delegate_search(ctx, queries=["climate change"]) + + # 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"], + ) + orchestrator.analysis_agent.run = AsyncMock(return_value=mock_result) + + # 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 + 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", + ) + orchestrator.synthesis_agent.run = AsyncMock(return_value=mock_result) + + # Create context and call the tool + ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage()) + result = await generate_report(ctx) + + # Verify the synthesis agent was called + orchestrator.synthesis_agent.run.assert_called_once() + + # Verify we got a ResearchReport + assert isinstance(result, ResearchReport) + assert result.title == "Climate Change Research"