Adapt clarification agent, basic orchestrator

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.analysis_agent import AnalysisAgent, AnalysisResult
-from haiku.rag.research.dependencies import ResearchDependencies
+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",
+]

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."""

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

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"