diff --git a/src/haiku/rag/research/dependencies.py b/src/haiku/rag/research/dependencies.py index dc74abee..fa90eedb 100644 --- a/src/haiku/rag/research/dependencies.py +++ b/src/haiku/rag/research/dependencies.py @@ -20,6 +20,9 @@ class ResearchContext(BaseModel): search_results: list[dict[str, Any]] = Field( default_factory=list, description="Accumulated search results" ) + qa_responses: list[dict[str, Any]] = Field( + default_factory=list, description="Question-answer pairs with sources" + ) insights: list[str] = Field( default_factory=list, description="Key insights discovered" ) @@ -39,6 +42,18 @@ class ResearchContext(BaseModel): } ) + def add_qa_response( + self, question: str, answer: str, sources: list["SearchResult"] + ) -> None: + """Add a QA response with its source documents.""" + self.qa_responses.append( + { + "question": question, + "answer": answer, + "sources": sources, + } + ) + def add_insight(self, insight: str) -> None: """Add a key insight.""" if insight not in self.insights: diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index 6792e93e..f73db6ec 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -202,30 +202,40 @@ class ResearchOrchestrator(BaseResearchAgent): 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)}" + search_prompt = ", ".join(search_target) elif iteration < len(context.sub_questions): # Use pre-planned sub-questions - search_prompt = f"Search for: {context.sub_questions[iteration]}" + search_prompt = context.sub_questions[iteration] else: - # Fall back to original question with variation - search_prompt = f"Additional search for: {question}" + # Fall back to original question + search_prompt = question # Search phase - directly call the search agent if console: - console.print(f"\n[bold cyan]🔍 Searching:[/bold cyan] {search_prompt}") + console.print( + f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {search_prompt}" + ) await self.search_agent.run(search_prompt, deps=deps) - if console and context.search_results: - latest_results = context.search_results[-1] - console.print( - f" Found [green]{len(latest_results.get('results', []))} documents[/green]" - ) - for i, result in enumerate(latest_results.get("results", [])[:3], 1): + if console: + # Show documents found + if context.search_results: + latest_results = context.search_results[-1] console.print( - f" {i}. Score: [yellow]{result.score:.3f}[/yellow] - {result.document_uri}" + 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: if console: diff --git a/src/haiku/rag/research/search_agent.py b/src/haiku/rag/research/search_agent.py index 80b52977..d06b3830 100644 --- a/src/haiku/rag/research/search_agent.py +++ b/src/haiku/rag/research/search_agent.py @@ -1,47 +1,69 @@ -"""Search specialist agent for document retrieval.""" +"""Search specialist agent that answers questions using RAG.""" + +from typing import Any from pydantic_ai import RunContext 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 document search and retrieval.""" + """Agent specialized in answering questions using RAG search.""" def __init__(self, provider: str, model: str): - # No specific output type needed - the tool handles everything - super().__init__(provider, model) + # 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: + """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: + # Get the sources from the last search (which the tool just stored) + if deps.context.search_results: + last_search = deps.context.search_results[-1] + sources = last_search.get("results", []) + deps.context.add_qa_response(prompt, result.output, sources) + + return result def get_system_prompt(self) -> str: - 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. Understand the search query and context - 2. Execute targeted searches to find relevant documents + return """You are a search and question-answering specialist. - Use the search tool to perform the searches on the knowledge base.""" + 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.""" def register_tools(self) -> None: """Register search-specific tools.""" @self.agent.tool - async def search( + async def search_and_answer( ctx: RunContext[ResearchDependencies], query: str, limit: int = 5, - ) -> list[tuple[Chunk, float]]: - """Execute search and return raw results from client.""" + ) -> str: + """Search for information and provide context for answering the question.""" # Use the default hybrid search search_results = await ctx.deps.client.search(query, limit=limit) # Expand context for better relevance expanded = await ctx.deps.client.expand_context(search_results) - # Store in context (convert to SearchResult for context storage) + # Convert to SearchResult for context storage from haiku.rag.research.base import SearchResult results_for_context = [] + context_texts = [] + for chunk, score in expanded: results_for_context.append( SearchResult( @@ -51,7 +73,16 @@ class SearchSpecialistAgent(BaseResearchAgent): metadata={"chunk_id": chunk.id} if chunk.id else {}, ) ) + context_texts.append(chunk.content) + + # Store raw search results for analysis by other agents ctx.deps.context.add_search_result(query, results_for_context) - # Return raw chunk, score tuples - return expanded + # Format context for the LLM to answer the question + if context_texts: + context = "\n\n---\n\n".join(context_texts) + return f"Based on the following information from the knowledge base:\n\n{context}\n\nAnswer the question: {query}" + else: + return ( + f"No relevant information found in the knowledge base for: {query}" + ) diff --git a/tests/research/test_search_agent.py b/tests/research/test_search_agent.py index 96661353..6105c9f6 100644 --- a/tests/research/test_search_agent.py +++ b/tests/research/test_search_agent.py @@ -77,7 +77,7 @@ async def test_search_agent_has_search_tool(): tools = test_model.last_model_request_parameters.function_tools assert tools is not None assert len(tools) == 1 - assert tools[0].name == "search" + assert tools[0].name == "search_and_answer" @pytest.mark.asyncio @@ -96,21 +96,19 @@ async def test_search_single_query(mock_client, research_deps): agent = SearchSpecialistAgent(provider="openai", model="gpt-4") # Get the search tool - search_tool = get_agent_tool(agent, "search") + search_tool = get_agent_tool(agent, "search_and_answer") assert search_tool is not None # Test the tool ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage()) - results = await search_tool(ctx, query="climate change") + result = await search_tool(ctx, query="climate change") - # Verify results - should be list of (Chunk, float) tuples - assert isinstance(results, list) - assert len(results) == 2 - 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 result - should be a formatted string with context + assert isinstance(result, str) + assert "Climate change is a global phenomenon" in result + assert "Rising temperatures affect ecosystems" in result - # Verify mock was called + # Verify mock was called with default limit mock_client.search.assert_called_once_with("climate change", limit=5) mock_client.expand_context.assert_called_once() @@ -133,18 +131,18 @@ async def test_search_with_limit(mock_client, research_deps): agent = SearchSpecialistAgent(provider="openai", model="gpt-4") # Get the search tool - search_tool = get_agent_tool(agent, "search") + search_tool = get_agent_tool(agent, "search_and_answer") assert search_tool is not None # Test the tool with limit ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage()) - results = await search_tool(ctx, query="test query", limit=3) + result = await search_tool(ctx, query="test query", limit=3) - # 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 result is a formatted string + assert isinstance(result, str) + assert "Content 1" in result + assert "Content 2" in result + assert "Content 3" in result # Verify mock was called with correct limit mock_client.search.assert_called_once_with("test query", limit=3) @@ -161,7 +159,7 @@ async def test_search_updates_context(mock_client, research_deps): agent = SearchSpecialistAgent(provider="openai", model="gpt-4") # Get the search tool - search_tool = get_agent_tool(agent, "search") + search_tool = get_agent_tool(agent, "search_and_answer") assert search_tool is not None # Test the tool