diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index 127d421d..9f7c4d4b 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -22,18 +22,18 @@ class ResearchPlan(BaseModel): 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 + self, provider: str | None = Config.RESEARCH_PROVIDER, model: str | None = None ): + # Use provided values or fall back to config defaults + provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER + model = model or Config.RESEARCH_MODEL or Config.QA_MODEL + super().__init__(provider, model, output_type=ResearchPlan) self.search_agent = SearchSpecialistAgent(provider, model) @@ -55,7 +55,8 @@ class ResearchOrchestrator(BaseResearchAgent): - Breaks down complex questions into manageable parts - Identifies multiple search strategies - Defines clear success criteria - - Ensures thorough investigation""" + - Ensures thorough investigation + /no_think""" def register_tools(self) -> None: """Register orchestration tools.""" @@ -63,13 +64,21 @@ class ResearchOrchestrator(BaseResearchAgent): @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 + ) -> 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( @@ -158,10 +167,8 @@ class ResearchOrchestrator(BaseResearchAgent): 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 + assert 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): @@ -176,16 +183,18 @@ class ResearchOrchestrator(BaseResearchAgent): else: # Fall back to original question with variation search_prompt = f"Additional search for: {question}" + # Search phase - directly call the search agent - # Search phase - await self.run(search_prompt, deps=deps) + await self.search_agent.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) + await self.analysis_agent.run( + "Analyze the gathered information", deps=deps + ) # Clarification phase - evaluate completeness - clarification_result = await self.run( + 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.", @@ -202,7 +211,9 @@ class ResearchOrchestrator(BaseResearchAgent): break # Generate final report - report_result = await self.run("Generate the final research report", deps=deps) + report_result = await self.synthesis_agent.run( + "Generate the final research report", deps=deps + ) return ( report_result.output if hasattr(report_result, "output") else report_result ) diff --git a/src/haiku/rag/research/search_agent.py b/src/haiku/rag/research/search_agent.py index 2a341cbd..80b52977 100644 --- a/src/haiku/rag/research/search_agent.py +++ b/src/haiku/rag/research/search_agent.py @@ -1,26 +1,26 @@ -"""Search specialist agent for advanced document retrieval.""" +"""Search specialist agent for document retrieval.""" from pydantic_ai import RunContext -from haiku.rag.research.base import BaseResearchAgent, SearchResult +from haiku.rag.research.base import BaseResearchAgent from haiku.rag.research.dependencies import ResearchDependencies +from haiku.rag.store.models.chunk import Chunk class SearchSpecialistAgent(BaseResearchAgent): - """Agent specialized in advanced document search and retrieval.""" + """Agent specialized in document search and retrieval.""" def __init__(self, provider: str, model: str): - super().__init__(provider, model, output_type=list[SearchResult]) + # No specific output type needed - the tool handles everything + super().__init__(provider, model) def get_system_prompt(self) -> str: - return """You are a search specialist agent focused on document retrieval. + return """You are a search specialist agent focused on document retrieval from a knowledge base that uses hybrid (semantic and full-text search) search. Your role is to: - 1. Generate multiple search queries from different perspectives - 2. Identify key terms and synonyms for comprehensive search - 3. Execute searches and rank results by relevance - 4. Return the most relevant documents for the research question + 1. Understand the search query and context + 2. Execute targeted searches to find relevant documents - Use the search tools to explore the knowledge base thoroughly.""" + Use the search tool to perform the searches on the knowledge base.""" def register_tools(self) -> None: """Register search-specific tools.""" @@ -28,46 +28,30 @@ class SearchSpecialistAgent(BaseResearchAgent): @self.agent.tool async def search( ctx: RunContext[ResearchDependencies], - queries: str | list[str], + query: str, limit: int = 5, - ) -> list[SearchResult]: - """Execute search with single or multiple query variants.""" - # Normalize to list - query_list = [queries] if isinstance(queries, str) else queries + ) -> list[tuple[Chunk, float]]: + """Execute search and return raw results from client.""" + # Use the default hybrid search + search_results = await ctx.deps.client.search(query, limit=limit) - all_results = [] - seen_chunk_ids = set() + # Expand context for better relevance + expanded = await ctx.deps.client.expand_context(search_results) - for query in query_list: - # Use the default hybrid search - search_results = await ctx.deps.client.search(query, limit=limit) + # Store in context (convert to SearchResult for context storage) + from haiku.rag.research.base import SearchResult - # Expand context for better relevance - expanded = await ctx.deps.client.expand_context(search_results) + results_for_context = [] + for chunk, score in expanded: + results_for_context.append( + SearchResult( + content=chunk.content, + score=score, + document_uri=chunk.document_uri or "", + metadata={"chunk_id": chunk.id} if chunk.id else {}, + ) + ) + ctx.deps.context.add_search_result(query, results_for_context) - for chunk, score in expanded: - # Avoid duplicates based on chunk ID - if chunk.id and chunk.id not in seen_chunk_ids: - seen_chunk_ids.add(chunk.id) - all_results.append( - SearchResult( - content=chunk.content, - score=score, - document_uri=chunk.document_uri or "", - metadata={"query": query, "chunk_id": chunk.id}, - ) - ) - - # Sort by score and limit results - all_results.sort(key=lambda x: x.score, reverse=True) - final_results = all_results[: limit * len(query_list)] - - # Store in context - query_summary = ( - query_list[0] - if len(query_list) == 1 - else f"Multi-query: {len(query_list)} variants" - ) - ctx.deps.context.add_search_result(query_summary, final_results) - - return final_results + # Return raw chunk, score tuples + return expanded diff --git a/tests/research/test_search_agent.py b/tests/research/test_search_agent.py index 435cb6c8..96661353 100644 --- a/tests/research/test_search_agent.py +++ b/tests/research/test_search_agent.py @@ -101,12 +101,14 @@ async def test_search_single_query(mock_client, research_deps): # Test the tool ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage()) - results = await search_tool(ctx, queries="climate change") + results = await search_tool(ctx, query="climate change") - # Verify results + # Verify results - should be list of (Chunk, float) tuples + assert isinstance(results, list) assert len(results) == 2 - assert results[0].content == "Climate change is a global phenomenon" - assert results[0].metadata["chunk_id"] == "chunk1" + assert results[0][0].content == "Climate change is a global phenomenon" + assert results[0][0].id == "chunk1" + assert results[0][1] == 0.8 # score # Verify mock was called mock_client.search.assert_called_once_with("climate change", limit=5) @@ -114,20 +116,19 @@ async def test_search_single_query(mock_client, research_deps): @pytest.mark.asyncio -async def test_search_multiple_queries_deduplication(mock_client, research_deps): - """Test deduplication when searching with multiple queries.""" - # Create chunks with duplicate IDs across queries - chunks_q1 = [ +async def test_search_with_limit(mock_client, research_deps): + """Test that search respects the limit parameter.""" + # Create more chunks than the limit + mock_chunks = [ create_mock_chunk("chunk1", "Content 1", 0.9), - create_mock_chunk("chunk2", "Content 2", 0.7), - ] - chunks_q2 = [ - create_mock_chunk("chunk1", "Content 1", 0.9), # Duplicate - create_mock_chunk("chunk3", "Content 3", 0.8), + create_mock_chunk("chunk2", "Content 2", 0.8), + create_mock_chunk("chunk3", "Content 3", 0.7), + create_mock_chunk("chunk4", "Content 4", 0.6), + create_mock_chunk("chunk5", "Content 5", 0.5), ] - mock_client.search.side_effect = [[chunks_q1[0]], [chunks_q2[0]]] - mock_client.expand_context.side_effect = [chunks_q1, chunks_q2] + mock_client.search.return_value = mock_chunks[:3] + mock_client.expand_context.return_value = mock_chunks[:3] agent = SearchSpecialistAgent(provider="openai", model="gpt-4") @@ -135,20 +136,18 @@ async def test_search_multiple_queries_deduplication(mock_client, research_deps) search_tool = get_agent_tool(agent, "search") assert search_tool is not None - # Test the tool + # Test the tool with limit ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage()) - results = await search_tool(ctx, queries=["query1", "query2"]) + results = await search_tool(ctx, query="test query", limit=3) - # Check deduplication - chunk1 should appear only once - chunk_ids = [r.metadata["chunk_id"] for r in results] - assert chunk_ids.count("chunk1") == 1 - assert "chunk2" in chunk_ids - assert "chunk3" in chunk_ids + # Verify results respect limit + assert isinstance(results, list) + assert len(results) == 3 + assert all(isinstance(r, tuple) and len(r) == 2 for r in results) + assert all(isinstance(r[0], Chunk) and isinstance(r[1], float) for r in results) - # Verify sorting by score - assert all( - results[i].score >= results[i + 1].score for i in range(len(results) - 1) - ) + # Verify mock was called with correct limit + mock_client.search.assert_called_once_with("test query", limit=3) @pytest.mark.asyncio @@ -167,7 +166,7 @@ async def test_search_updates_context(mock_client, research_deps): # Test the tool ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage()) - await search_tool(ctx, queries="test query") + await search_tool(ctx, query="test query") # Verify context was updated assert len(research_deps.context.search_results) == 1